python-roborock 2.39.2__tar.gz → 2.40.1__tar.gz

{python_roborock-2.39.2 → python_roborock-2.40.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: python-roborock
-Version: 2.39.2
+Version: 2.40.1
 Summary: A package to control Roborock vacuums.
 Home-page: https://github.com/humbertogontijo/python-roborock
 License: GPL-3.0-only

{python_roborock-2.39.2 → python_roborock-2.40.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "python-roborock"
-version = "2.39.2"
+version = "2.40.1"
 description = "A package to control Roborock vacuums."
 authors = ["humbertogontijo <humbertogontijo@users.noreply.github.com>"]
 license = "GPL-3.0-only"

python_roborock-2.40.1/roborock/broadcast_protocol.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import json
+import logging
+from asyncio import BaseTransport, Lock
+
+from construct import (  # type: ignore
+    Bytes,
+    Checksum,
+    GreedyBytes,
+    Int16ub,
+    Int32ub,
+    Prefixed,
+    RawCopy,
+    Struct,
+)
+from Crypto.Cipher import AES
+
+from roborock import RoborockException
+from roborock.containers import BroadcastMessage
+from roborock.protocol import EncryptionAdapter, Utils, _Parser
+
+_LOGGER = logging.getLogger(__name__)
+
+BROADCAST_TOKEN = b"qWKYcdQWrbm9hPqe"
+
+
+class RoborockProtocol(asyncio.DatagramProtocol):
+    def __init__(self, timeout: int = 5):
+        self.timeout = timeout
+        self.transport: BaseTransport | None = None
+        self.devices_found: list[BroadcastMessage] = []
+        self._mutex = Lock()
+
+    def datagram_received(self, data: bytes, _):
+        """Handle incoming broadcast datagrams."""
+        try:
+            version = data[:3]
+            if version == b"L01":
+                [parsed_msg], _ = L01Parser.parse(data)
+                encrypted_payload = parsed_msg.payload
+                if encrypted_payload is None:
+                    raise RoborockException("No encrypted payload found in broadcast message")
+                ciphertext = encrypted_payload[:-16]
+                tag = encrypted_payload[-16:]
+
+                key = hashlib.sha256(BROADCAST_TOKEN).digest()
+                iv_digest_input = data[:9]
+                digest = hashlib.sha256(iv_digest_input).digest()
+                iv = digest[:12]
+
+                cipher = AES.new(key, AES.MODE_GCM, nonce=iv)
+                decrypted_payload_bytes = cipher.decrypt_and_verify(ciphertext, tag)
+                json_payload = json.loads(decrypted_payload_bytes)
+                parsed_message = BroadcastMessage(duid=json_payload["duid"], ip=json_payload["ip"], version=version)
+                _LOGGER.debug(f"Received L01 broadcast: {parsed_message}")
+                self.devices_found.append(parsed_message)
+            else:
+                # Fallback to the original protocol parser for other versions
+                [broadcast_message], _ = BroadcastParser.parse(data)
+                if broadcast_message.payload:
+                    json_payload = json.loads(broadcast_message.payload)
+                    parsed_message = BroadcastMessage(duid=json_payload["duid"], ip=json_payload["ip"], version=version)
+                    _LOGGER.debug(f"Received broadcast: {parsed_message}")
+                    self.devices_found.append(parsed_message)
+        except Exception as e:
+            _LOGGER.warning(f"Failed to decode message: {data!r}. Error: {e}")
+
+    async def discover(self) -> list[BroadcastMessage]:
+        async with self._mutex:
+            try:
+                loop = asyncio.get_event_loop()
+                self.transport, _ = await loop.create_datagram_endpoint(lambda: self, local_addr=("0.0.0.0", 58866))
+                await asyncio.sleep(self.timeout)
+                return self.devices_found
+            finally:
+                self.close()
+                self.devices_found = []
+
+    def close(self):
+        self.transport.close() if self.transport else None
+
+
+_BroadcastMessage = Struct(
+    "message"
+    / RawCopy(
+        Struct(
+            "version" / Bytes(3),
+            "seq" / Int32ub,
+            "protocol" / Int16ub,
+            "payload" / EncryptionAdapter(lambda ctx: BROADCAST_TOKEN),
+        )
+    ),
+    "checksum" / Checksum(Int32ub, Utils.crc, lambda ctx: ctx.message.data),
+)
+
+_L01BroadcastMessage = Struct(
+    "message"
+    / RawCopy(
+        Struct(
+            "version" / Bytes(3),
+            "field1" / Bytes(4),  # Unknown field
+            "field2" / Bytes(2),  # Unknown field
+            "payload" / Prefixed(Int16ub, GreedyBytes),  # Encrypted payload with length prefix
+        )
+    ),
+    "checksum" / Checksum(Int32ub, Utils.crc, lambda ctx: ctx.message.data),
+)
+
+
+BroadcastParser: _Parser = _Parser(_BroadcastMessage, False)
+L01Parser: _Parser = _Parser(_L01BroadcastMessage, False)

{python_roborock-2.39.2 → python_roborock-2.40.1}/roborock/containers.py

@@ -783,6 +783,7 @@ class FlowLedStatus(RoborockBase):
 class BroadcastMessage(RoborockBase):
     duid: str
     ip: str
+    version: bytes
 
 
 class ServerTimer(NamedTuple):

{python_roborock-2.39.2 → python_roborock-2.40.1}/roborock/protocol.py

@@ -304,10 +304,10 @@ class _Parser:
             messages.append(
                 RoborockMessage(
                     version=message.message.value.version,
-                    seq=message.message.value.seq,
+                    seq=message.message.value.get("seq"),
                     random=message.message.value.get("random"),
                     timestamp=message.message.value.get("timestamp"),
-                    protocol=message.message.value.protocol,
+                    protocol=message.message.value.get("protocol"),
                     payload=message.message.value.payload,
                 )
             )

{python_roborock-2.39.2 → python_roborock-2.40.1}/roborock/protocols/v1_protocol.py

@@ -157,19 +157,18 @@ class MapResponse:
     """The map data, decrypted and decompressed."""
 
 
-def create_map_response_decoder(security_data: SecurityData) -> Callable[[RoborockMessage], MapResponse]:
+def create_map_response_decoder(security_data: SecurityData) -> Callable[[RoborockMessage], MapResponse | None]:
     """Create a decoder for V1 map response messages."""
 
-    def _decode_map_response(message: RoborockMessage) -> MapResponse:
+    def _decode_map_response(message: RoborockMessage) -> MapResponse | None:
         """Decode a V1 map response message."""
         if not message.payload or len(message.payload) < 24:
             raise RoborockException("Invalid V1 map response format: missing payload")
         header, body = message.payload[:24], message.payload[24:]
         [endpoint, _, request_id, _] = struct.unpack("<8s8sH6s", header)
         if not endpoint.decode().startswith(security_data.endpoint):
-            raise RoborockException(
-                f"Invalid V1 map response endpoint: {endpoint!r}, expected {security_data.endpoint!r}"
-            )
+            _LOGGER.debug("Received map response requested not made by this device, ignoring.")
+            return None
         try:
             decrypted = Utils.decrypt_cbc(body, security_data.nonce)
         except ValueError as err:

{python_roborock-2.39.2 → python_roborock-2.40.1}/roborock/version_1_apis/roborock_client_v1.py

@@ -150,7 +150,7 @@ class RoborockClientV1(RoborockClient, ABC):
     """Roborock client base class for version 1 devices."""
 
     _listeners: dict[str, ListenerModel] = {}
-    _map_response_decoder: Callable[[RoborockMessage], MapResponse] | None = None
+    _map_response_decoder: Callable[[RoborockMessage], MapResponse | None] | None = None
 
     def __init__(self, device_info: DeviceData, security_data: SecurityData | None) -> None:
         """Initializes the Roborock client."""
@@ -439,13 +439,14 @@ class RoborockClientV1(RoborockClient, ABC):
                 elif data.payload and protocol == RoborockMessageProtocol.MAP_RESPONSE:
                     if self._map_response_decoder is not None:
                         map_response = self._map_response_decoder(data)
-                        queue = self._waiting_queue.get(map_response.request_id)
-                        if queue:
-                            queue.set_result(map_response.data)
-                        else:
-                            self._logger.debug(
-                                "Received unsolicited map response for request_id %s", map_response.request_id
-                            )
+                        if map_response is not None:
+                            queue = self._waiting_queue.get(map_response.request_id)
+                            if queue:
+                                queue.set_result(map_response.data)
+                            else:
+                                self._logger.debug(
+                                    "Received unsolicited map response for request_id %s", map_response.request_id
+                                )
                 else:
                     queue = self._waiting_queue.get(data.seq)
                     if queue:

python_roborock-2.39.2/roborock/broadcast_protocol.py

@@ -1,68 +0,0 @@
-from __future__ import annotations
-
-import asyncio
-import json
-import logging
-from asyncio import BaseTransport, Lock
-
-from construct import (  # type: ignore
-    Bytes,
-    Checksum,
-    Int16ub,
-    Int32ub,
-    RawCopy,
-    Struct,
-)
-
-from roborock.containers import BroadcastMessage
-from roborock.protocol import EncryptionAdapter, Utils, _Parser
-
-_LOGGER = logging.getLogger(__name__)
-
-BROADCAST_TOKEN = b"qWKYcdQWrbm9hPqe"
-
-
-class RoborockProtocol(asyncio.DatagramProtocol):
-    def __init__(self, timeout: int = 5):
-        self.timeout = timeout
-        self.transport: BaseTransport | None = None
-        self.devices_found: list[BroadcastMessage] = []
-        self._mutex = Lock()
-
-    def datagram_received(self, data, _):
-        [broadcast_message], _ = BroadcastParser.parse(data)
-        if broadcast_message.payload:
-            parsed_message = BroadcastMessage.from_dict(json.loads(broadcast_message.payload))
-            _LOGGER.debug(f"Received broadcast: {parsed_message}")
-            self.devices_found.append(parsed_message)
-
-    async def discover(self):
-        async with self._mutex:
-            try:
-                loop = asyncio.get_event_loop()
-                self.transport, _ = await loop.create_datagram_endpoint(lambda: self, local_addr=("0.0.0.0", 58866))
-                await asyncio.sleep(self.timeout)
-                return self.devices_found
-            finally:
-                self.close()
-                self.devices_found = []
-
-    def close(self):
-        self.transport.close() if self.transport else None
-
-
-_BroadcastMessage = Struct(
-    "message"
-    / RawCopy(
-        Struct(
-            "version" / Bytes(3),
-            "seq" / Int32ub,
-            "protocol" / Int16ub,
-            "payload" / EncryptionAdapter(lambda ctx: BROADCAST_TOKEN),
-        )
-    ),
-    "checksum" / Checksum(Int32ub, Utils.crc, lambda ctx: ctx.message.data),
-)
-
-
-BroadcastParser: _Parser = _Parser(_BroadcastMessage, False)