chore: using python construct for data parsing (#48)

* chore: using python construct for data parsing

* chore: linting

* chore: linting

* fix: roborock message protocol

* fix: change local api constructor

* chore: linting

* chore: linting

Author: humbertogontijo

.pre-commit-config.yaml

@@ -40,7 +40,7 @@ repos:
                 args:
                     - --fix
     -   repo: https://github.com/pre-commit/mirrors-mypy
-        rev: v0.931
+        rev: v1.2.0
         hooks:
             -   id: mypy
                 exclude: cli.py

pyproject.toml

@@ -21,7 +21,7 @@ roborock = "roborock.cli:main"
 [tool.poetry.dependencies]
 python = "^3.9"
 click = ">=8"
-aiohttp = "*"
+aiohttp = "^3.8.2"
 async-timeout = "*"
 pycryptodome = "~3.17.0"
 pycryptodomex = {version = "~3.17.0", markers = "sys_platform == 'darwin'"}
@@ -37,12 +37,13 @@ build-backend = "poetry.core.masonry.api"
 [tool.poetry.dev-dependencies]
 pytest-asyncio = "*"
 pytest = "*"
-pre-commit = "*"
+pre-commit = "^3.2.2"
 mypy = "*"
 ruff = "*"
 isort = "*"
 black = "*"
 codespell = "*"
+pyshark = "^0.6"
 
 [tool.semantic_release]
 branch = "main"

roborock/api.py

@@ -4,7 +4,6 @@
 
 import asyncio
 import base64
-import gzip
 import hashlib
 import hmac
 import json
@@ -17,8 +16,6 @@
 from typing import Any, Callable, Coroutine, Optional
 
 import aiohttp
-from Crypto.Cipher import AES
-from Crypto.Util.Padding import unpad
 
 from .code_mappings import RoborockDockTypeCode
 from .containers import (
@@ -46,6 +43,7 @@
     RoborockUrlException,
     VacuumError,
 )
+from .protocol import Utils
 from .roborock_future import RoborockFuture
 from .roborock_message import RoborockMessage
 from .roborock_typing import DeviceProp, DockSummary, RoborockCommand
@@ -148,15 +146,13 @@ def on_message_received(self, messages: list[RoborockMessage]) -> None:
                     payload = data.payload[0:24]
                     [endpoint, _, request_id, _] = struct.unpack("<15sBH6s", payload)
                     if endpoint.decode().startswith(self._endpoint):
-                        iv = bytes(AES.block_size)
-                        decipher = AES.new(self._nonce, AES.MODE_CBC, iv)
-                        decrypted = unpad(decipher.decrypt(data.payload[24:]), AES.block_size)
-                        decrypted = gzip.decompress(decrypted)
+                        decrypted = Utils.decrypt_cbc(data.payload[24:], self._nonce)
+                        decompressed = Utils.decompress(decrypted)
                         queue = self._waiting_queue.get(request_id)
                         if queue:
-                            if isinstance(decrypted, list):
-                                decrypted = decrypted[0]
-                            queue.resolve((decrypted, None))
+                            if isinstance(decompressed, list):
+                                decompressed = decompressed[0]
+                            queue.resolve((decompressed, None))
                 else:
                     queue = self._waiting_queue.get(data.seq)
                     if queue:

roborock/cli.py

@@ -123,9 +123,18 @@ async def command(ctx, cmd, device_id, params):
     home_data = login_data.home_data
     devices = home_data.devices + home_data.received_devices
     device = next((device for device in devices if device.duid == device_id), None)
+    if device is None:
+        raise RoborockException("No device found")
     model_specification = next(
-        (product.model_specification for product in home_data.products if product.did == device.duid), None
+        (
+            product.model_specification
+            for product in home_data.products
+            if device is not None and product.did == device.duid
+        ),
+        None,
     )
+    if model_specification is None:
+        raise RoborockException(f"Could not find model specifications for device {device.name}")
     device_info = RoborockDeviceInfo(device=device, model_specification=model_specification)
     mqtt_client = RoborockMqttClient(login_data.user_data, device_info)
     await mqtt_client.send_command(cmd, params)

roborock/cloud_api.py

@@ -13,9 +13,9 @@
 from .api import COMMANDS_SECURED, KEEPALIVE, RoborockClient, md5hex
 from .containers import RoborockDeviceInfo, UserData
 from .exceptions import CommandVacuumError, RoborockException, VacuumError
-from .protocol import Utils
+from .protocol import MessageParser, Utils
 from .roborock_future import RoborockFuture
-from .roborock_message import RoborockMessage, RoborockParser
+from .roborock_message import RoborockMessage
 from .roborock_typing import RoborockCommand
 
 _LOGGER = logging.getLogger(__name__)
@@ -77,7 +77,7 @@ def on_connect(self, *args, **kwargs):
     def on_message(self, *args, **kwargs):
         _, __, msg = args
         try:
-            messages, _ = RoborockParser.decode(msg.payload, self.device_info.device.local_key)
+            messages, _ = MessageParser.parse(msg.payload, local_key=self.device_info.device.local_key)
             super().on_message_received(messages)
         except Exception as ex:
             _LOGGER.exception(ex)
@@ -156,7 +156,7 @@ async def send_command(self, method: RoborockCommand, params: Optional[list | di
         response_protocol = 301 if method in COMMANDS_SECURED else 102
         roborock_message = RoborockMessage(timestamp=timestamp, protocol=request_protocol, payload=payload)
         local_key = self.device_info.device.local_key
-        msg = RoborockParser.encode(roborock_message, local_key)
+        msg = MessageParser.build(roborock_message, local_key)
         self._send_msg_raw(msg)
         (response, err) = await self._async_response(request_id, response_protocol)
         if err:

roborock/containers.py

@@ -388,11 +388,6 @@ class RoborockDeviceInfo(RoborockBase):
     model_specification: ModelSpecification
 
 
-@dataclass
-class RoborockLocalDeviceInfo(RoborockDeviceInfo):
-    network_info: NetworkInfo
-
-
 @dataclass
 class RoomMapping(RoborockBase):
     segment_id: int

roborock/local_api.py

@@ -8,20 +8,21 @@
 import async_timeout
 
 from .api import COMMANDS_SECURED, QUEUE_TIMEOUT, RoborockClient
-from .containers import RoborockLocalDeviceInfo
+from .containers import RoborockDeviceInfo
 from .exceptions import CommandVacuumError, RoborockConnectionException, RoborockException
-from .roborock_message import AP_CONFIG, RoborockMessage, RoborockParser
+from .protocol import AP_CONFIG, MessageParser
+from .roborock_message import RoborockMessage
 from .roborock_typing import CommandInfoMap, RoborockCommand
 from .util import get_running_loop_or_create_one
 
 _LOGGER = logging.getLogger(__name__)
 
 
 class RoborockLocalClient(RoborockClient, asyncio.Protocol):
-    def __init__(self, device_info: RoborockLocalDeviceInfo):
+    def __init__(self, device_info: RoborockDeviceInfo, ip: str):
         super().__init__("abc", device_info)
         self.loop = get_running_loop_or_create_one()
-        self.ip = device_info.network_info.ip
+        self.ip = ip
         self._batch_structs: list[RoborockMessage] = []
         self._executing = False
         self.remaining = b""
@@ -32,8 +33,7 @@ def data_received(self, message):
         if self.remaining:
             message = self.remaining + message
             self.remaining = b""
-        (parser_msg, remaining) = RoborockParser.decode(message, self.device_info.device.local_key)
-        self.remaining = remaining
+        parser_msg, self.remaining = MessageParser.parse(message, local_key=self.device_info.device.local_key)
         self.on_message_received(parser_msg)
 
     def connection_lost(self, exc: Optional[Exception]):
@@ -92,11 +92,16 @@ async def send_command(self, method: RoborockCommand, params: Optional[list | di
         return (await self.send_message(roborock_message))[0]
 
     async def async_local_response(self, roborock_message: RoborockMessage):
-        request_id = roborock_message.get_request_id()
-        if request_id is None:
+        method = roborock_message.get_method()
+        request_id: int | None
+        if method and not method.startswith("get"):
             request_id = roborock_message.seq
-        # response_protocol = 5 if roborock_message.prefix == secured_prefix else 4
-        response_protocol = 4
+            response_protocol = 5
+        else:
+            request_id = roborock_message.get_request_id()
+            response_protocol = 4
+        if request_id is None:
+            raise RoborockException(f"Failed build message {roborock_message}")
         (response, err) = await self._async_response(request_id, response_protocol)
         if err:
             raise CommandVacuumError("", err) from err
@@ -116,7 +121,7 @@ async def send_message(self, roborock_messages: list[RoborockMessage] | Roborock
         if isinstance(roborock_messages, RoborockMessage):
             roborock_messages = [roborock_messages]
         local_key = self.device_info.device.local_key
-        msg = RoborockParser.encode(roborock_messages, local_key)
+        msg = MessageParser.build(roborock_messages, local_key=local_key)
         # Send the command to the Roborock device
         _LOGGER.debug(f"Requesting device with {roborock_messages}")
         self._send_msg_raw(msg)