Split KLAP IoT updates into smaller requests and validate KLAP responses (#1829)

Some IoT devices respond over KLAP with a plain `200 OK` html page
instead of a proper response when the response would be too large to
construct, causing a "not a multiple of the block length" decryption
error (e.g., KS200M in #1828).

This PR lowers the maximum response size to 2048 bytes for IoT devices
using KLAP, so that the update is split into multiple smaller requests
like it is already done for bulbs.

In details:
- `IotDevice.max_device_response_size` returns 2048 over KLAP. Bulbs
  keep 4096 and XOR devices keep 16 KiB.
- `KlapTransport` checks the response length and signature before
  decrypting, and raises a `KasaException` telling what was wrong
  instead of the decryption error.

Fixes #1828

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

kasa/iot/iotdevice.py

@@ -24,7 +24,7 @@ from warnings import warn
 
 from ..device import Device, DeviceInfo, WifiNetwork
 from ..device_type import DeviceType
-from ..deviceconfig import DeviceConfig
+from ..deviceconfig import DeviceConfig, DeviceEncryptionType
 from ..exceptions import KasaException, TimeoutError, UnsupportedDeviceError
 from ..feature import Feature
 from ..module import Module
@@ -730,6 +730,11 @@ class IotDevice(Device):
     @property
     def max_device_response_size(self) -> int:
         """Returns the maximum response size the device can safely construct."""
+        # Over KLAP, some devices answer a request whose response would be
+        # larger than about 4 KiB with a plain html page instead of a response,
+        # and have even less room right after a state change.
+        if self.config.connection_type.encryption_type is DeviceEncryptionType.Klap:
+            return 2048
         return 16 * 1024
 
     @property

kasa/transports/klaptransport.py

@@ -46,6 +46,7 @@ import asyncio
 import base64
 import datetime
 import hashlib
+import hmac
 import logging
 import secrets
 import struct
@@ -396,6 +397,18 @@ class KlapTransport(BaseTransport):
             if TYPE_CHECKING:
                 assert self._encryption_session
                 assert isinstance(response_data, bytes)
+            # A message is a 32 byte signature followed by whole AES blocks.
+            # Some devices answer with a plain html page instead.
+            if len(response_data) < 48 or (len(response_data) - 32) % 16:
+                raise KasaException(
+                    f"Device {self._host} sent a response that is not a KLAP "
+                    f"message ({len(response_data)} bytes): {response_data[:64]!r}"
+                )
+            if not self._encryption_session.is_signed(response_data, seq):
+                raise KasaException(
+                    f"Device {self._host} sent a response with an invalid "
+                    f"signature for request with seq {seq}"
+                )
             try:
                 decrypted_response = self._encryption_session.decrypt(response_data)
             except Exception as ex:
@@ -535,6 +548,11 @@ class KlapEncryptionSession:
         ).digest()
         return (signature + ciphertext, self._seq)
 
+    def is_signed(self, msg: bytes, seq: int) -> bool:
+        """Return whether the message is signed for the request with seq."""
+        signature = _sha256(self._sig + PACK_SIGNED_LONG(seq) + msg[32:])
+        return hmac.compare_digest(msg[:32], signature)
+
     def decrypt(self, msg: bytes) -> str:
         """Decrypt the data."""
         decryptor = self._cipher.decryptor()

tests/transports/test_klaptransport.py

@@ -11,11 +11,17 @@ from contextlib import nullcontext as does_not_raise
 
 import aiohttp
 import pytest
+from cryptography.hazmat.primitives.ciphers import Cipher, modes
 from pytest_mock import MockerFixture
 from yarl import URL
 
 from kasa.credentials import DEFAULT_CREDENTIALS, Credentials, get_default_credentials
-from kasa.deviceconfig import DeviceConfig
+from kasa.deviceconfig import (
+    DeviceConfig,
+    DeviceConnectionParameters,
+    DeviceEncryptionType,
+    DeviceFamily,
+)
 from kasa.exceptions import (
     AuthenticationError,
     KasaException,
@@ -24,17 +30,27 @@ from kasa.exceptions import (
     _RetryableError,
 )
 from kasa.httpclient import HttpClient
+from kasa.iot import IotBulb, IotDevice, IotPlug, IotWallSwitch
 from kasa.protocols import IotProtocol, SmartProtocol
 from kasa.transports.aestransport import AesTransport
 from kasa.transports.klaptransport import (
+    PACK_SIGNED_LONG,
     KlapEncryptionSession,
     KlapTransport,
     KlapTransportV2,
     _sha256,
 )
 
+from ..device_fixtures import device_iot, parametrize
+
 DUMMY_QUERY = {"foobar": {"foo": "bar", "bar": "foo"}}
 
+KLAP_CONNECTION = DeviceConnectionParameters(
+    DeviceFamily.IotSmartPlugSwitch, DeviceEncryptionType.Klap, login_version=2
+)
+
+ks200m_iot = parametrize("ks200m", model_filter={"KS200M"}, protocol_filter={"IOT"})
+
 # Transport tests are not designed for real devices
 pytestmark = [pytest.mark.requires_dummy]
 
@@ -209,7 +225,8 @@ async def test_protocol_logging(
     logging.getLogger("kasa").setLevel(log_level)
 
     def _return_encrypted(*_, **__):
-        # Do the encrypt just before returning the value so the incrementing sequence number is correct
+        # Devices sign the response for the sequence number of the request
+        encryption_session._seq -= 1
         encrypted, seq = encryption_session.encrypt('{"great":"success"}')
         return 200, encrypted
 
@@ -292,36 +309,63 @@ async def test_transport_decrypt(mocker: MockerFixture) -> None:
     assert d == resp
 
 
-async def test_transport_decrypt_error(
+def _logged_in_transport() -> KlapTransport:
-    mocker: MockerFixture, caplog: pytest.LogCaptureFixture
-) -> None:
-    """Test that a decryption error raises a kasa exception."""
-    d = {"great": "success"}
-
     seed = secrets.token_bytes(16)
     auth_hash = KlapTransport.generate_auth_hash(Credentials("foo", "bar"))
-    encryption_session = KlapEncryptionSession(seed, seed, auth_hash)
 
     transport = KlapTransport(config=DeviceConfig(host="127.0.0.1"))
     transport._handshake_done = True
     transport._session_expire_at = time.monotonic() + 60
-    transport._encryption_session = encryption_session
+    transport._encryption_session = KlapEncryptionSession(seed, seed, auth_hash)
+    return transport
+
+
+@pytest.mark.parametrize("mismatch", ["key", "seq"])
+async def test_transport_signature_error(mocker: MockerFixture, mismatch: str) -> None:
+    """Test that a response not signed for the request raises a kasa exception."""
+    transport = _logged_in_transport()
 
     async def _return_response(url: URL, params=None, data=None, *_, **__):
         assert transport._encryption_session is not None
         encryption_session = KlapEncryptionSession(
-            secrets.token_bytes(16),
+            secrets.token_bytes(16)
+            if mismatch == "key"
+            else transport._encryption_session.local_seed,
             transport._encryption_session.remote_seed,
             transport._encryption_session.user_hash,
         )
         seq = params.get("seq")
-        encryption_session._seq = seq - 1
+        encryption_session._seq = seq - 1 if mismatch == "key" else seq
-        encrypted, seq = encryption_session.encrypt(json.dumps(d))
+        encrypted, _seq = encryption_session.encrypt(json.dumps({"great": "success"}))
-        seq = seq
         return 200, encrypted
 
     mocker.patch.object(HttpClient, "post", side_effect=_return_response)
 
+    with pytest.raises(
+        KasaException,
+        match=re.escape(
+            "Device 127.0.0.1 sent a response with an invalid signature "
+            "for request with seq"
+        ),
+    ):
+        await transport.send(json.dumps({}))
+
+
+async def test_transport_decrypt_error(mocker: MockerFixture) -> None:
+    """Test that a signed response that fails to decrypt raises a kasa exception."""
+    transport = _logged_in_transport()
+
+    async def _return_response(url: URL, params=None, data=None, *_, **__):
+        session = transport._encryption_session
+        assert session is not None
+        seq = PACK_SIGNED_LONG(params.get("seq"))
+        # A block that decrypts to zeros, which is not valid PKCS7 padding
+        encryptor = Cipher(session._aes, modes.CBC(session._iv + seq)).encryptor()
+        ciphertext = encryptor.update(bytes(16)) + encryptor.finalize()
+        return 200, _sha256(session._sig + seq + ciphertext) + ciphertext
+
+    mocker.patch.object(HttpClient, "post", side_effect=_return_response)
+
     with pytest.raises(
         KasaException,
         match=re.escape("Error trying to decrypt device 127.0.0.1 response:"),
@@ -329,6 +373,90 @@ async def test_transport_decrypt_error(
         await transport.send(json.dumps({}))
 
 
+HTML_PAGE = b"<html><body><center>200 OK</center></body></html>"
+
+
+@pytest.mark.parametrize(
+    "response",
+    [
+        pytest.param(HTML_PAGE, id="html-page"),
+        pytest.param(b"", id="empty"),
+        pytest.param(secrets.token_bytes(32), id="signature-only"),
+        pytest.param(secrets.token_bytes(32 + 17), id="partial-block"),
+    ],
+)
+async def test_transport_invalid_response(
+    mocker: MockerFixture, response: bytes
+) -> None:
+    """Test that a response that is not a KLAP message raises a kasa exception."""
+    transport = _logged_in_transport()
+    mocker.patch.object(HttpClient, "post", return_value=(200, response))
+
+    with pytest.raises(
+        KasaException,
+        match=re.escape(
+            "Device 127.0.0.1 sent a response that is not a KLAP message "
+            f"({len(response)} bytes): {response[:64]!r}"
+        ),
+    ):
+        await transport.send(json.dumps({}))
+
+
+@pytest.mark.parametrize("transport_class", [KlapTransport, KlapTransportV2])
+async def test_protocol_invalid_response(
+    mocker: MockerFixture, transport_class: type[KlapTransport]
+) -> None:
+    """Test that an invalid response is not retried and resets the session."""
+    server_seed = secrets.token_bytes(16)
+    client_credentials = Credentials("foo", "bar")
+    device_auth_hash = transport_class.generate_auth_hash(client_credentials)
+    handshakes = 0
+    requests = 0
+
+    async def _return_response(url: URL, params=None, data=None, *_, **__):
+        nonlocal handshakes, requests
+
+        if url == URL("http://127.0.0.1:80/app/handshake1"):
+            handshakes += 1
+            seed_auth_hash = _sha256(
+                data + server_seed + device_auth_hash
+                if transport_class is KlapTransportV2
+                else data + device_auth_hash
+            )
+            return _mock_response(200, server_seed + seed_auth_hash)
+        elif url == URL("http://127.0.0.1:80/app/handshake2"):
+            return _mock_response(200, b"")
+        elif url == URL("http://127.0.0.1:80/app/request"):
+            requests += 1
+            if requests == 1:
+                return _mock_response(200, HTML_PAGE)
+            assert transport._encryption_session is not None
+            encryption_session = KlapEncryptionSession(
+                transport._encryption_session.local_seed,
+                transport._encryption_session.remote_seed,
+                transport._encryption_session.user_hash,
+            )
+            encryption_session._seq = params.get("seq") - 1
+            encrypted, _seq = encryption_session.encrypt('{"great": "success"}')
+            return _mock_response(200, encrypted)
+
+    mocker.patch.object(aiohttp.ClientSession, "post", side_effect=_return_response)
+
+    config = DeviceConfig("127.0.0.1", credentials=client_credentials)
+    transport = transport_class(config=config)
+    protocol = IotProtocol(transport=transport)
+
+    with pytest.raises(KasaException, match="is not a KLAP message"):
+        await protocol.query(DUMMY_QUERY)
+    assert requests == 1
+    assert handshakes == 1
+    assert transport._handshake_done is False
+
+    assert await protocol.query(DUMMY_QUERY) == {"great": "success"}
+    assert requests == 2
+    assert handshakes == 2
+
+
 @pytest.mark.parametrize(
     ("device_credentials", "expectation"),
     [
@@ -594,3 +722,46 @@ async def test_port_override() -> None:
     transport = KlapTransport(config=config)
 
     assert str(transport._app_url) == "http://127.0.0.1:12345/app"
+
+
+@device_iot
+async def test_max_device_response_size_klap(
+    dev: IotDevice, mocker: MockerFixture
+) -> None:
+    """Test that devices other than bulbs use a smaller budget over KLAP."""
+    is_bulb = isinstance(dev, IotBulb)
+    assert dev.max_device_response_size == (4096 if is_bulb else 16 * 1024)
+
+    mocker.patch.object(dev.config, "connection_type", KLAP_CONNECTION)
+    expected = 4096 if is_bulb else 2048
+    assert dev.max_device_response_size == expected
+    for child in dev.children:
+        assert child.max_device_response_size == expected
+
+
+def test_max_device_response_size_before_update() -> None:
+    """Test that the budget follows the connection settings before an update."""
+    assert IotPlug("127.0.0.1").max_device_response_size == 16 * 1024
+
+    config = DeviceConfig("127.0.0.1", connection_type=KLAP_CONNECTION)
+    assert IotPlug("127.0.0.1", config=config).max_device_response_size == 2048
+
+
+@ks200m_iot
+async def test_ks200m_update_split_klap(
+    dev: IotWallSwitch, mocker: MockerFixture
+) -> None:
+    """Check that a KS200M update over KLAP is split into smaller requests."""
+    spy = mocker.spy(dev.protocol, "query")
+    await dev.update()
+    assert spy.call_count == 1
+
+    mocker.patch.object(dev.config, "connection_type", KLAP_CONNECTION)
+    spy.reset_mock()
+    await dev.update()
+
+    assert [sorted(call.args[0]) for call in spy.call_args_list] == [
+        ["schedule", "system"],
+        ["schedule"],
+        ["anti_theft", "cnCloud", "smartlife.iot.LAS", "smartlife.iot.PIR", "time"],
+    ]