python-kasa · sdb9696 · Jan 23, 2024 · Jan 23, 2024 · Jan 23, 2024 · Jan 23, 2024
diff --git a/kasa/aestransport.py b/kasa/aestransport.py
@@ -8,7 +8,7 @@
 import hashlib
 import logging
 import time
-from typing import Dict, Optional, cast
+from typing import TYPE_CHECKING, AsyncGenerator, Dict, Optional, cast
 
 from cryptography.hazmat.primitives import padding, serialization
 from cryptography.hazmat.primitives.asymmetric import padding as asymmetric_padding
@@ -55,6 +55,8 @@ class AesTransport(BaseTransport):
         "requestByApp": "true",
         "Accept": "application/json",
     }
+    CONTENT_LENGTH = "Content-Length"
+    KEY_PAIR_CONTENT_LENGTH = 314
 
     def __init__(
         self,
@@ -86,6 +88,8 @@ def __init__(
 
         self._login_token = None
 
+        self._key_pair: Optional[KeyPair] = None
+
         _LOGGER.debug("Created AES transport for %s", self._host)
 
     @property
@@ -185,34 +189,44 @@ async def perform_login(self):
         self._handle_response_error_code(resp_dict, "Error logging in")
         self._login_token = resp_dict["result"]["token"]
 
-    async def perform_handshake(self):
-        """Perform the handshake."""
-        _LOGGER.debug("Will perform handshaking...")
-        _LOGGER.debug("Generating keypair")
-
-        self._handshake_done = False
-        self._session_expire_at = None
-        self._session_cookie = None
-
-        url = f"http://{self._host}/app"
-        key_pair = KeyPair.create_key_pair()
+    async def _generate_key_pair_payload(self) -> AsyncGenerator:
+        """Generate the request body and return an ascyn_generator.
 
+        This prevents the key pair being generated unless a connection
+        can be made to the device.
+        """
+        _LOGGER.debug("Generating keypair")
+        self._key_pair = KeyPair.create_key_pair()
         pub_key = (
             "-----BEGIN PUBLIC KEY-----\n"
-            + key_pair.get_public_key()
+            + self._key_pair.get_public_key()  # type: ignore[union-attr]
             + "\n-----END PUBLIC KEY-----\n"
         )
         handshake_params = {"key": pub_key}
         _LOGGER.debug(f"Handshake params: {handshake_params}")
-
         request_body = {"method": "handshake", "params": handshake_params}
-
         _LOGGER.debug(f"Request {request_body}")
+        yield json_dumps(request_body).encode()
 
+    async def perform_handshake(self):
+        """Perform the handshake."""
+        _LOGGER.debug("Will perform handshaking...")
+
+        self._key_pair = None
+        self._handshake_done = False
+        self._session_expire_at = None
+        self._session_cookie = None
+
+        url = f"http://{self._host}/app"
+        # Device needs the content length or it will response with 500
+        headers = {
+            **self.COMMON_HEADERS,
+            self.CONTENT_LENGTH: str(self.KEY_PAIR_CONTENT_LENGTH),
+        }
         status_code, resp_dict = await self._http_client.post(
             url,
-            json=request_body,
-            headers=self.COMMON_HEADERS,
+            json=self._generate_key_pair_payload(),
+            headers=headers,
             cookies_dict=self._session_cookie,
         )

@@ -240,8 +254,10 @@ async def perform_handshake(self):
             self._session_cookie = {self.SESSION_COOKIE_NAME: cookie}
 
         self._session_expire_at = time.time() + 86400
+        if TYPE_CHECKING:
+            assert self._key_pair is not None  # pragma: no cover
         self._encryption_session = AesEncyptionSession.create_from_keypair(
-            handshake_key, key_pair
+            handshake_key, self._key_pair
         )
 
         self._handshake_done = True

diff --git a/kasa/httpclient.py b/kasa/httpclient.py
@@ -41,14 +41,25 @@ async def post(
         *,
         params: Optional[Dict[str, Any]] = None,
         data: Optional[bytes] = None,
-        json: Optional[Dict] = None,
+        json: Optional[Union[Dict, Any]] = None,
         headers: Optional[Dict[str, str]] = None,
         cookies_dict: Optional[Dict[str, str]] = None,
     ) -> Tuple[int, Optional[Union[Dict, bytes]]]:
-        """Send an http post request to the device."""
+        """Send an http post request to the device.
+
+        If the request is provided via the json parameter json will be returned.
+        """
         response_data = None
         self._last_url = url
         self.client.cookie_jar.clear()
+        return_json = bool(json)
+        # If json is not a dict send as data.
+        # This allows the json parameter to be used to pass other
+        # types of data such as async_generator and still have json
+        # returned.
+        if json and not isinstance(json, Dict):
+            data = json
+            json = None
         try:
             resp = await self.client.post(
                 url,
@@ -62,7 +73,7 @@ async def post(
             async with resp:
                 if resp.status == 200:
                     response_data = await resp.read()
-                    if json:
+                    if return_json:
                         response_data = json_loads(response_data.decode())
 
         except (aiohttp.ServerDisconnectedError, aiohttp.ClientOSError) as ex:

diff --git a/kasa/json.py b/kasa/json.py
@@ -11,5 +11,9 @@ def dumps(obj, *, default=None):
 except ImportError:
     import json
 
-    dumps = json.dumps
+    def dumps(obj, *, default=None):
+        """Dump JSON."""
+        # Separators specified for consistency with orjson
+        return json.dumps(obj, separators=(",", ":"))
+
     loads = json.loads
diff --git a/kasa/tests/test_aestransport.py b/kasa/tests/test_aestransport.py
@@ -169,7 +169,10 @@ def __init__(self, host, status_code=200, error_code=0, inner_error_code=0):
         self.inner_error_code = inner_error_code
         self.http_client = HttpClient(DeviceConfig(self.host))
 
-    async def post(self, url, params=None, json=None, *_, **__):
+    async def post(self, url, params=None, json=None, data=None, *_, **__):
+        if data:
+            async for item in data:
+                json = json_loads(item.decode())
         return await self._post(url, json)
 
     async def _post(self, url, json):