feat: dynamic calculated prefixes

Author: humbertogontijo

roborock/cli.py

@@ -85,7 +85,8 @@ async def _discover(ctx):
         raise Exception("You need to login first")
     client = RoborockApiClient(login_data.email)
     home_data = await client.get_home_data(login_data.user_data)
-    context.update(LoginData(**login_data.as_dict(), home_data=home_data))
+    login_data.home_data = home_data
+    context.update(login_data)
     click.echo(f"Discovered devices {', '.join([device.name for device in home_data.get_all_devices()])}")
 
 
@@ -128,7 +129,7 @@ async def command(ctx, cmd, device_id, params):
     devices = home_data.devices + home_data.received_devices
     device = next(device for device in devices if device.duid == device_id)
     model = next(
-        (product.model for product in home_data.products if device is not None and product.did == device.duid),
+        (product.model for product in home_data.products if device is not None and product.id == device.product_id),
         None,
     )
     if model is None:

roborock/cloud_api.py

@@ -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 = MessageParser.build(roborock_message, local_key)
+        msg = MessageParser.build(roborock_message, local_key, False)
         self._send_msg_raw(msg)
         (response, err) = await self._async_response(request_id, response_protocol)
         if err:

roborock/local_api.py

@@ -60,7 +60,7 @@ async def async_connect(self) -> None:
                 raise RoborockConnectionException(f"Failed connecting to {self.host}") from e
 
     def sync_disconnect(self) -> None:
-        if self.transport:
+        if self.transport and self.loop.is_running():
             self.transport.close()
 
     async def async_disconnect(self) -> None:
@@ -77,18 +77,15 @@ def build_roborock_message(self, method: RoborockCommand, params: Optional[list
         command = CommandInfoMap.get(method)
         if command is None:
             raise RoborockException(f"No prefix found for {method}")
-        prefix = command.prefix
         request_protocol = 4
         return RoborockMessage(
-            prefix=prefix,
             timestamp=timestamp,
             protocol=request_protocol,
             payload=payload,
         )
 
     async def ping(self):
-        command_info = CommandInfoMap[RoborockCommand.NONE]
-        roborock_message = RoborockMessage(prefix=command_info.prefix, protocol=AP_CONFIG, payload=b"")
+        roborock_message = RoborockMessage(protocol=AP_CONFIG, payload=b"")
         return (await self.send_message(roborock_message))[0]
 
     async def send_command(self, method: RoborockCommand, params: Optional[list | dict] = None):

roborock/protocol.py

@@ -25,6 +25,8 @@
     RawCopy,
     Struct,
     bytestringtype,
+    stream_seek,
+    stream_tell,
 )
 from Crypto.Cipher import AES
 from Crypto.Util.Padding import pad, unpad
@@ -223,43 +225,51 @@ def _parse(self, stream, context, path):
         return hash1
 
 
-class OptionalPrefix(Construct):
+class PrefixedStruct(Struct):
     def _parse(self, stream, context, path):
         subcon1 = Peek(Optional(Const(b"1.0")))
         peek_version = subcon1.parse_stream(stream, **context)
         if peek_version is None:
             subcon2 = Bytes(4)
-            return subcon2.parse_stream(stream, **context)
-        return b""
+            subcon2.parse_stream(stream, **context)
+        return super()._parse(stream, context, path)
 
     def _build(self, obj, stream, context, path):
-        if obj is not None:
-            subcon1 = Bytes(4)
-            subcon1.build_stream(obj, stream, **context)
+        prefixed = context.search("prefixed")
+        if not prefixed:
+            return super()._build(obj, stream, context, path)
+        offset = stream_tell(stream, path)
+        stream_seek(stream, offset + 4, 0, path)
+        super()._build(obj, stream, context, path)
+        new_offset = stream_tell(stream, path)
+        subcon1 = Bytes(4)
+        stream_seek(stream, offset, 0, path)
+        subcon1.build_stream(new_offset - offset - subcon1.sizeof(**context), stream, **context)
+        stream_seek(stream, new_offset + 4, 0, path)
         return obj
 
 
+_Message = RawCopy(
+    Struct(
+        "version" / Const(b"1.0"),
+        "seq" / Int32ub,
+        "random" / Int32ub,
+        "timestamp" / Int32ub,
+        "protocol" / Int16ub,
+        "payload"
+        / EncryptionAdapter(
+            lambda ctx: Utils.md5(
+                Utils.encode_timestamp(ctx.timestamp) + Utils.ensure_bytes(ctx.search("local_key")) + SALT
+            ),
+        ),
+    )
+)
+
 _Messages = Struct(
     "messages"
     / GreedyRange(
-        Struct(
-            "prefix" / OptionalPrefix(),
-            "message"
-            / RawCopy(
-                Struct(
-                    "version" / Const(b"1.0"),
-                    "seq" / Int32ub,
-                    "random" / Int32ub,
-                    "timestamp" / Int32ub,
-                    "protocol" / Int16ub,
-                    "payload"
-                    / EncryptionAdapter(
-                        lambda ctx: Utils.md5(
-                            Utils.encode_timestamp(ctx.timestamp) + Utils.ensure_bytes(ctx.search("local_key")) + SALT
-                        ),
-                    ),
-                )
-            ),
+        PrefixedStruct(
+            "message" / _Message,
             "checksum" / OptionalChecksum(Optional(Int32ub), Utils.crc, lambda ctx: ctx.message.data),
         )
     ),
@@ -294,7 +304,6 @@ def parse(self, data: bytes, local_key: str | None = None) -> tuple[list[Roboroc
         for message in parsed_messages:
             messages.append(
                 RoborockMessage(
-                    prefix=message.get("prefix"),
                     version=message.message.value.version,
                     seq=message.message.value.seq,
                     random=message.message.value.get("random"),
@@ -306,14 +315,15 @@ def parse(self, data: bytes, local_key: str | None = None) -> tuple[list[Roboroc
         remaining = parsed.get("remaining") or b""
         return messages, remaining
 
-    def build(self, roborock_messages: list[RoborockMessage] | RoborockMessage, local_key: str) -> bytes:
+    def build(
+        self, roborock_messages: list[RoborockMessage] | RoborockMessage, local_key: str, prefixed: bool = True
+    ) -> bytes:
         if isinstance(roborock_messages, RoborockMessage):
             roborock_messages = [roborock_messages]
         messages = []
         for roborock_message in roborock_messages:
             messages.append(
                 {
-                    "prefix": roborock_message.prefix,
                     "message": {
                         "value": {
                             "version": roborock_message.version,
@@ -326,7 +336,7 @@ def build(self, roborock_messages: list[RoborockMessage] | RoborockMessage, loca
                     },
                 }
             )
-        return self.con.build({"messages": [message for message in messages]}, local_key=local_key)
+        return self.con.build({"messages": [message for message in messages]}, local_key=local_key, prefixed=prefixed)
 
 
 MessageParser: _Parser = _Parser(_Messages, True)

roborock/roborock_message.py

@@ -5,7 +5,6 @@
 import time
 from dataclasses import dataclass
 from random import randint
-from typing import Optional
 
 from .roborock_typing import RoborockCommand
 
@@ -15,7 +14,6 @@ class RoborockMessage:
     protocol: int
     payload: bytes
     seq: int = randint(100000, 999999)
-    prefix: Optional[bytes] = None
     version: bytes = b"1.0"
     random: int = randint(10000, 99999)
     timestamp: int = math.floor(time.time())