Merge branch 'main' into mqtt_error_handling

Author: Lash-L

.github/workflows/ci.yml

@@ -25,7 +25,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v5
-      - uses: actions/setup-python@v5
+      - uses: actions/setup-python@v6
         with:
           python-version: "3.11"
       - uses: pre-commit/action@v3.0.1
@@ -42,7 +42,7 @@ jobs:
     steps:
       - uses: actions/checkout@v5
       - name: Set up Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: ${{ matrix.python-version }}
       - uses: snok/install-poetry@v1.4.1
@@ -75,12 +75,12 @@ jobs:
         persist-credentials: false
     - name: Python Semantic Release
       id: release
-      uses: python-semantic-release/python-semantic-release@v10.3.0
+      uses: python-semantic-release/python-semantic-release@v10.4.1
       with:
         github_token: ${{ secrets.GH_TOKEN }}
 
     - name: Publish package distributions to PyPI
-      uses: pypa/gh-action-pypi-publish@v1.12.4
+      uses: pypa/gh-action-pypi-publish@v1.13.0
 
       # NOTE: DO NOT wrap the conditional in ${{ }} as it will always evaluate to true.
       # See https://github.com/actions/runner/issues/1173

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "python-roborock"
-version = "2.39.0"
+version = "2.45.0"
 description = "A package to control Roborock vacuums."
 authors = ["humbertogontijo <humbertogontijo@users.noreply.github.com>"]
 license = "GPL-3.0-only"
@@ -32,6 +32,7 @@ construct = "^2.10.57"
 vacuum-map-parser-roborock = "*"
 pyrate-limiter = "^3.7.0"
 aiomqtt = "^2.3.2"
+click-shell = "^2.1"
 
 
 [build-system]

pyproject.toml

@@ -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)