diff --git a/kasa/iot/iotdevice.py b/kasa/iot/iotdevice.py index e5e93fc5..ba3f0a6b 100755 --- a/kasa/iot/iotdevice.py +++ b/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 diff --git a/kasa/transports/klaptransport.py b/kasa/transports/klaptransport.py index b9fbfa87..e0ca936f 100644 --- a/kasa/transports/klaptransport.py +++ b/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() diff --git a/tests/transports/test_klaptransport.py b/tests/transports/test_klaptransport.py index b1931b99..8dc886fa 100644 --- a/tests/transports/test_klaptransport.py +++ b/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( - mocker: MockerFixture, caplog: pytest.LogCaptureFixture -) -> None: - """Test that a decryption error raises a kasa exception.""" - d = {"great": "success"} - +def _logged_in_transport() -> KlapTransport: 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 - encrypted, seq = encryption_session.encrypt(json.dumps(d)) - seq = seq + encryption_session._seq = seq - 1 if mismatch == "key" else seq + encrypted, _seq = encryption_session.encrypt(json.dumps({"great": "success"})) 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"
200 OK
" + + +@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"], + ]