Enable newer encrypted discovery protocol (#1168)

kasa/discover.py

@@ -82,13 +82,16 @@ Discovering a single device returns a kasa.Device object.
 from __future__ import annotations
 
 import asyncio
+import base64
 import binascii
 import ipaddress
 import logging
+import secrets
 import socket
+import struct
 from collections.abc import Awaitable
 from pprint import pformat as pf
-from typing import Any, Callable, Dict, Optional, Type, cast
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Type, cast
 
 # When support for cpython older than 3.11 is dropped
 # async_timeout can be replaced with asyncio.timeout
@@ -96,6 +99,7 @@ from async_timeout import timeout as asyncio_timeout
 from pydantic.v1 import BaseModel, ValidationError
 
 from kasa import Device
+from kasa.aestransport import AesEncyptionSession, KeyPair
 from kasa.credentials import Credentials
 from kasa.device_factory import (
     get_device_class_from_family,
@@ -133,6 +137,46 @@ NEW_DISCOVERY_REDACTORS: dict[str, Callable[[Any], Any] | None] = {
 }
 
 
+class _AesDiscoveryQuery:
+    keypair: KeyPair | None = None
+
+    @classmethod
+    def generate_query(cls):
+        if not cls.keypair:
+            cls.keypair = KeyPair.create_key_pair(key_size=2048)
+        secret = secrets.token_bytes(4)
+
+        key_payload = {"params": {"rsa_key": cls.keypair.get_public_pem().decode()}}
+
+        key_payload_bytes = json_dumps(key_payload).encode()
+        # https://labs.withsecure.com/advisories/tp-link-ac1750-pwn2own-2019
+        version = 2  # version of tdp
+        msg_type = 0
+        op_code = 1  # probe
+        msg_size = len(key_payload_bytes)
+        flags = 17
+        padding_byte = 0  # blank byte
+        device_serial = int.from_bytes(secret, "big")
+        initial_crc = 0x5A6B7C8D
+
+        disco_header = struct.pack(
+            ">BBHHBBII",
+            version,
+            msg_type,
+            op_code,
+            msg_size,
+            flags,
+            padding_byte,
+            device_serial,
+            initial_crc,
+        )
+
+        query = bytearray(disco_header + key_payload_bytes)
+        crc = binascii.crc32(query).to_bytes(length=4, byteorder="big")
+        query[12:16] = crc
+        return query
+
+
 class _DiscoverProtocol(asyncio.DatagramProtocol):
     """Implementation of the discovery protocol handler.
 
@@ -224,15 +268,21 @@ class _DiscoverProtocol(asyncio.DatagramProtocol):
         _LOGGER.debug("[DISCOVERY] %s >> %s", self.target, Discover.DISCOVERY_QUERY)
         encrypted_req = XorEncryption.encrypt(req)
         sleep_between_packets = self.discovery_timeout / self.discovery_packets
+
+        aes_discovery_query = _AesDiscoveryQuery.generate_query()
         for _ in range(self.discovery_packets):
             if self.target in self.seen_hosts:  # Stop sending for discover_single
                 break
             self.transport.sendto(encrypted_req[4:], self.target_1)  # type: ignore
             self.transport.sendto(Discover.DISCOVERY_QUERY_2, self.target_2)  # type: ignore
+            self.transport.sendto(aes_discovery_query, self.target_2)  # type: ignore
             await asyncio.sleep(sleep_between_packets)
 
     def datagram_received(self, data, addr) -> None:
         """Handle discovery responses."""
+        if TYPE_CHECKING:
+            assert _AesDiscoveryQuery.keypair
+
         ip, port = addr
         # Prevent multiple entries due multiple broadcasts
         if ip in self.seen_hosts:
@@ -395,7 +445,8 @@ class Discover:
         credentials: Credentials | None = None,
         username: str | None = None,
         password: str | None = None,
-    ) -> Device:
+        on_unsupported: OnUnsupportedCallable | None = None,
+    ) -> Device | None:
         """Discover a single device by the given IP address.
 
         It is generally preferred to avoid :func:`discover_single()` and
@@ -465,7 +516,11 @@ class Discover:
             dev.host = host
             return dev
         elif ip in protocol.unsupported_device_exceptions:
-            raise protocol.unsupported_device_exceptions[ip]
+            if on_unsupported:
+                await on_unsupported(protocol.unsupported_device_exceptions[ip])
+                return None
+            else:
+                raise protocol.unsupported_device_exceptions[ip]
         elif ip in protocol.invalid_device_exceptions:
             raise protocol.invalid_device_exceptions[ip]
         else:
@@ -512,6 +567,25 @@ class Discover:
         device.update_from_discover_info(info)
         return device
 
+    @staticmethod
+    def _decrypt_discovery_data(discovery_result: DiscoveryResult) -> None:
+        if TYPE_CHECKING:
+            assert discovery_result.encrypt_info
+            assert _AesDiscoveryQuery.keypair
+        encryped_key = discovery_result.encrypt_info.key
+        encrypted_data = discovery_result.encrypt_info.data
+
+        key_and_iv = _AesDiscoveryQuery.keypair.decrypt_discovery_key(
+            base64.b64decode(encryped_key.encode())
+        )
+
+        key, iv = key_and_iv[:16], key_and_iv[16:]
+
+        session = AesEncyptionSession(key, iv)
+        decrypted_data = session.decrypt(encrypted_data)
+
+        discovery_result.decrypted_data = json_loads(decrypted_data)
+
     @staticmethod
     def _get_device_instance(
         data: bytes,
@@ -528,6 +602,8 @@ class Discover:
             ) from ex
         try:
             discovery_result = DiscoveryResult(**info["result"])
+            if discovery_result.encrypt_info:
+                Discover._decrypt_discovery_data(discovery_result)
         except ValidationError as ex:
             if debug_enabled:
                 data = (
@@ -547,9 +623,19 @@ class Discover:
         type_ = discovery_result.device_type
 
         try:
+            if not (
+                encrypt_type := discovery_result.mgt_encrypt_schm.encrypt_type
+            ) and (encrypt_info := discovery_result.encrypt_info):
+                encrypt_type = encrypt_info.sym_schm
+            if not encrypt_type:
+                raise UnsupportedDeviceError(
+                    f"Unsupported device {config.host} of type {type_} "
+                    + "with no encryption type",
+                    discovery_result=discovery_result.get_dict(),
+                )
             config.connection_type = DeviceConnectionParameters.from_values(
                 type_,
-                discovery_result.mgt_encrypt_schm.encrypt_type,
+                encrypt_type,
                 discovery_result.mgt_encrypt_schm.lv,
             )
         except KasaException as ex:
@@ -593,21 +679,35 @@ class EncryptionScheme(BaseModel):
     """Base model for encryption scheme of discovery result."""
 
     is_support_https: bool
-    encrypt_type: str
-    http_port: int
+    encrypt_type: Optional[str]  # noqa: UP007
+    http_port: Optional[int] = None  # noqa: UP007
     lv: Optional[int] = None  # noqa: UP007
 
 
+class EncryptionInfo(BaseModel):
+    """Base model for encryption info of discovery result."""
+
+    sym_schm: str
+    key: str
+    data: str
+
+
 class DiscoveryResult(BaseModel):
     """Base model for discovery result."""
 
     device_type: str
     device_model: str
+    device_name: Optional[str]  # noqa: UP007
     ip: str
     mac: str
     mgt_encrypt_schm: EncryptionScheme
+    encrypt_info: Optional[EncryptionInfo] = None  # noqa: UP007
+    encrypt_type: Optional[list[str]] = None  # noqa: UP007
+    decrypted_data: Optional[dict] = None  # noqa: UP007
     device_id: str
 
+    firmware_version: Optional[str] = None  # noqa: UP007
+    hardware_version: Optional[str] = None  # noqa: UP007
     hw_ver: Optional[str] = None  # noqa: UP007
     owner: Optional[str] = None  # noqa: UP007
     is_support_iot_cloud: Optional[bool] = None  # noqa: UP007