Adding support for the new encryption protocol (updated version)

kasa/__init__.py

@@ -13,6 +13,7 @@
 """
 from importlib_metadata import version  # type: ignore
 
+from kasa.auth import Auth
 from kasa.discover import Discover
 from kasa.emeterstatus import EmeterStatus
 from kasa.exceptions import SmartDeviceException
@@ -28,6 +29,7 @@
 
 
 __all__ = [
+    "Auth",
     "Discover",
     "TPLinkSmartHomeProtocol",
     "SmartBulb",

kasa/auth.py

@@ -0,0 +1,21 @@
+"""Authentication class for KASA username / passwords."""
+from hashlib import md5
+
+
+class Auth:
+    """Authentication for Kasa KLAP authentication."""
+
+    def __init__(self, user: str = "", password: str = ""):
+        self.user = user
+        self.password = password
+        self.md5user = md5(user.encode()).digest()
+        self.md5password = md5(password.encode()).digest()
+        self.md5auth = md5(self.md5user + self.md5password).digest()
+
+    def authenticator(self):
+        """Return the KLAP authenticator for these credentials."""
+        return self.md5auth
+
+    def owner(self):
+        """Return the MD5 hash of the username in this object."""
+        return self.md5user

kasa/cli.py

@@ -6,6 +6,7 @@
 import asyncclick as click
 
 from kasa import (
+    Auth,
     Discover,
     SmartBulb,
     SmartDevice,
@@ -44,9 +45,24 @@
 @click.option("--plug", default=False, is_flag=True)
 @click.option("--lightstrip", default=False, is_flag=True)
 @click.option("--strip", default=False, is_flag=True)
+@click.option("--klap", default=False, is_flag=True)
+@click.option(
+    "--user",
+    default="",
+    required=False,
+    help="Username/email address to authenticate to device.",
+)
+@click.option(
+    "--password",
+    default="",
+    required=False,
+    help="Password to use to authenticate to device.",
+)
 @click.version_option()
 @click.pass_context
-async def cli(ctx, host, alias, target, debug, bulb, plug, lightstrip, strip):
+async def cli(
+    ctx, host, alias, target, debug, bulb, plug, lightstrip, strip, klap, user, password
+):
     """A tool for controlling TP-Link smart home devices."""  # noqa
     if debug:
         logging.basicConfig(level=logging.DEBUG)
@@ -65,18 +81,27 @@ async def cli(ctx, host, alias, target, debug, bulb, plug, lightstrip, strip):
             click.echo(f"No device with name {alias} found")
             return
 
+    if password != "" and user == "":
+        click.echo("Using a password requires a username")
+        return
+
+    if klap or user != "":
+        authentication = Auth(user=user, password=password)
+    else:
+        authentication = None
+
     if host is None:
         click.echo("No host name given, trying discovery..")
         await ctx.invoke(discover)
         return
     else:
         if not bulb and not plug and not strip and not lightstrip:
             click.echo("No --strip nor --bulb nor --plug given, discovering..")
-            dev = await Discover.discover_single(host)
+            dev = await Discover.discover_single(host, authentication)
         elif bulb:
             dev = SmartBulb(host)
         elif plug:
-            dev = SmartPlug(host)
+            dev = SmartPlug(host, authentication)
         elif strip:
             dev = SmartStrip(host)
         elif lightstrip:
@@ -135,8 +160,18 @@ async def join(dev: SmartDevice, ssid, password, keytype):
 async def discover(ctx, timeout, discover_only, dump_raw):
     """Discover devices in the network."""
     target = ctx.parent.params["target"]
-    click.echo(f"Discovering devices on {target} for {timeout} seconds")
-    found_devs = await Discover.discover(target=target, timeout=timeout)
+    user = ctx.parent.params["user"]
+    password = ctx.parent.params["password"]
+
+    if user:
+        auth = Auth(user=user, password=password)
+    else:
+        auth = None
+
+    click.echo(f"Discovering devices for {timeout} seconds")
+    found_devs = await Discover.discover(
+        target=target, timeout=timeout, return_raw=dump_raw, authentication=auth
+    )
     if not discover_only:
         for ip, dev in found_devs.items():
             if dump_raw:

kasa/discover.py

@@ -1,10 +1,14 @@
 """Discovery module for TP-Link Smart Home devices."""
 import asyncio
+import binascii
+import hashlib
 import json
 import logging
 import socket
-from typing import Awaitable, Callable, Dict, Optional, Type, cast
+from typing import Awaitable, Callable, Dict, Mapping, Optional, Type, Union, cast
 
+from kasa.auth import Auth
+from kasa.klapprotocol import TPLinkKLAP
 from kasa.protocol import TPLinkSmartHomeProtocol
 from kasa.smartbulb import SmartBulb
 from kasa.smartdevice import SmartDevice, SmartDeviceException
@@ -35,13 +39,17 @@ def __init__(
         target: str = "255.255.255.255",
         discovery_packets: int = 3,
         interface: Optional[str] = None,
+        authentication: Optional[Auth] = None,
     ):
         self.transport = None
         self.discovery_packets = discovery_packets
         self.interface = interface
         self.on_discovered = on_discovered
         self.target = (target, Discover.DISCOVERY_PORT)
+        self.new_target = (target, Discover.NEW_DISCOVERY_PORT)
         self.discovered_devices = {}
+        self.authentication = authentication
+        self.emptyUser = hashlib.md5().digest()
 
     def connection_made(self, transport) -> None:
         """Set socket options for broadcasting."""
@@ -61,26 +69,49 @@ def do_discover(self) -> None:
         req = json.dumps(Discover.DISCOVERY_QUERY)
         _LOGGER.debug("[DISCOVERY] %s >> %s", self.target, Discover.DISCOVERY_QUERY)
         encrypted_req = TPLinkSmartHomeProtocol.encrypt(req)
+        new_req = binascii.unhexlify("020000010000000000000000463cb5d3")
         for i in range(self.discovery_packets):
             self.transport.sendto(encrypted_req[4:], self.target)  # type: ignore
+            self.transport.sendto(new_req, self.new_target)  # type: ignore
 
     def datagram_received(self, data, addr) -> None:
         """Handle discovery responses."""
         ip, port = addr
         if ip in self.discovered_devices:
             return
 
-        info = json.loads(TPLinkSmartHomeProtocol.decrypt(data))
-        _LOGGER.debug("[DISCOVERY] %s << %s", ip, info)
-
-        try:
+        if port == 9999:
+            info = json.loads(TPLinkSmartHomeProtocol.decrypt(data))
             device_class = Discover._get_device_class(info)
-        except SmartDeviceException as ex:
-            _LOGGER.debug("Unable to find device type from %s: %s", info, ex)
-            return
+            device = device_class(ip)
+        else:
+            info = json.loads(data[16:])
+            device_class = Discover._get_new_device_class(info)
+            owner = Discover._get_new_owner(info)
+            if owner is not None:
+                owner_bin = bytes.fromhex(owner)
+
+            _LOGGER.debug(
+                "[DISCOVERY] Device owner is %s, empty owner is %s",
+                owner_bin,
+                self.emptyUser,
+            )
+            if owner is None or owner == "" or owner_bin == self.emptyUser:
+                _LOGGER.debug("[DISCOVERY] Device %s has no owner", ip)
+                device = device_class(ip, Auth())
+            elif (
+                self.authentication is not None
+                and owner_bin == self.authentication.owner()
+            ):
+                _LOGGER.debug("[DISCOVERY] Device %s has authenticated owner", ip)
+                device = device_class(ip, self.authentication)
+            else:
+                _LOGGER.debug("[DISCOVERY] Found %s with unknown owner %s", ip, owner)
+                return
+
+        _LOGGER.debug("[DISCOVERY] %s << %s", ip, info)
 
-        device = device_class(ip)
-        device.update_from_discover_info(info)
+        asyncio.ensure_future(device.update())
 
         self.discovered_devices[ip] = device
 
@@ -132,6 +163,8 @@ class Discover:
 
     DISCOVERY_PORT = 9999
 
+    NEW_DISCOVERY_PORT = 20002
+
     DISCOVERY_QUERY = {
         "system": {"get_sysinfo": None},
     }
@@ -144,7 +177,8 @@ async def discover(
         timeout=5,
         discovery_packets=3,
         interface=None,
-    ) -> DeviceDict:
+        authentication=None,
+    ) -> Mapping[str, Union[SmartDevice, Dict]]:
         """Discover supported devices.
 
         Sends discovery message to 255.255.255.255:9999 in order
@@ -171,6 +205,7 @@ async def discover(
                 on_discovered=on_discovered,
                 discovery_packets=discovery_packets,
                 interface=interface,
+                authentication=authentication,
             ),
             local_addr=("0.0.0.0", 0),
         )
@@ -187,20 +222,29 @@ async def discover(
         return protocol.discovered_devices
 
     @staticmethod
-    async def discover_single(host: str) -> SmartDevice:
+    async def discover_single(host: str, authentication=None) -> SmartDevice:
         """Discover a single device by the given IP address.
 
         :param host: Hostname of device to query
         :rtype: SmartDevice
         :return: Object for querying/controlling found device.
         """
-        protocol = TPLinkSmartHomeProtocol(host)
+        if authentication is None:
+            protocol = TPLinkSmartHomeProtocol(host)
+        else:
+            protocol = TPLinkKLAP(host, authentication)
+        # protocol = TPLinkSmartHomeProtocol(host)
 
         info = await protocol.query(Discover.DISCOVERY_QUERY)
 
         device_class = Discover._get_device_class(info)
-        dev = device_class(host)
-        await dev.update()
+        if device_class is not None:
+            if authentication is None:
+                dev = device_class(host)
+            else:
+                dev = device_class(host, authentication)
+            await dev.update()
+            return dev
 
         return dev
 
@@ -231,3 +275,43 @@ def _get_device_class(info: dict) -> Type[SmartDevice]:
             return SmartBulb
 
         raise SmartDeviceException("Unknown device type: %s" % type_)
+
+    @staticmethod
+    def _get_new_device_class(info: dict) -> Type[SmartDevice]:
+        """Find SmartDevice subclass given new discovery payload."""
+        if "result" not in info:
+            raise SmartDeviceException("No 'result' in discovery response")
+
+        if "device_type" not in info["result"]:
+            raise SmartDeviceException("No 'device_type' in discovery result")
+
+        dtype = info["result"]["device_type"]
+
+        if dtype == "IOT.SMARTPLUGSWITCH":
+            return SmartPlug
+
+        raise SmartDeviceException("Unknown device type: %s", dtype)
+
+    @staticmethod
+    def _get_new_owner(info: dict) -> Optional[str]:
+        """Find owner given new-style discovery payload."""
+        if "result" not in info:
+            raise SmartDeviceException("No 'result' in discovery response")
+
+        if "owner" not in info["result"]:
+            return None
+
+        return info["result"]["owner"]
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    loop = asyncio.get_event_loop()
+
+    async def _on_device(dev):
+        await dev.update()
+        _LOGGER.info("Got device: %s", dev)
+
+    devices = loop.run_until_complete(Discover.discover(on_discovered=_on_device))
+    for ip, dev in devices.items():
+        print(f"[{ip}] {dev}")