refactor: add A01 (#199)

* major: add A01

* chore: add init

* chore: fix commitlint?

* chore: fix commitlint

* chore: fix commitlint

* chore: change refactor to be major tag

* refactor: add A01

* feat: add a01

BREAKING CHANGE: You must now specify what version api you want to use with clients.

pyproject.toml

@@ -49,8 +49,20 @@ pyshark = "^0.6"
 branch = "main"
 version_toml = "pyproject.toml:tool.poetry.version"
 build_command = "pip install poetry && poetry build"
+[tool.semantic_release.commit_parser_options]
+allowed_tags = [
+    "chore",
+    "docs",
+    "feat",
+    "fix",
+    "refactor"
+]
+major_tags= ["refactor"]
 
 [tool.ruff]
 ignore = ["F403", "E741"]
 line-length = 120
 select=["E", "F", "UP", "I"]
+
+[tool.ruff.lint.per-file-ignores]
+"*/__init__.py" = ["F401"]

roborock/version_1_apis/__init__.py

@@ -0,0 +1,3 @@
+from .roborock_client_v1 import AttributeCache, RoborockClientV1
+from .roborock_local_client_v1 import RoborockLocalClientV1
+from .roborock_mqtt_client_v1 import RoborockMqttClientV1

roborock/version_a01_apis/__init__.py

@@ -0,0 +1,2 @@
+from .roborock_client_a01 import RoborockClientA01
+from .roborock_mqtt_client_a01 import RoborockMqttClientA01

roborock/version_a01_apis/roborock_client_a01.py

@@ -0,0 +1,100 @@
+import dataclasses
+import json
+import typing
+from collections.abc import Callable
+from datetime import time
+
+from Crypto.Cipher import AES
+from Crypto.Util.Padding import unpad
+
+from roborock import DeviceData
+from roborock.api import RoborockClient
+from roborock.code_mappings import (
+    DyadBrushSpeed,
+    DyadCleanMode,
+    DyadError,
+    DyadSelfCleanLevel,
+    DyadSelfCleanMode,
+    DyadSuction,
+    DyadWarmLevel,
+    DyadWaterLevel,
+    RoborockDyadStateCode,
+)
+from roborock.containers import DyadProductInfo, DyadSndState
+from roborock.roborock_message import (
+    RoborockDyadDataProtocol,
+    RoborockMessage,
+    RoborockMessageProtocol,
+)
+
+
+@dataclasses.dataclass
+class DyadProtocolCacheEntry:
+    post_process_fn: Callable
+    value: typing.Any | None = None
+
+
+# Right now this cache is not active, it was too much complexity for the initial addition of dyad.
+protocol_entries = {
+    RoborockDyadDataProtocol.STATUS: DyadProtocolCacheEntry(lambda val: RoborockDyadStateCode(val).name),
+    RoborockDyadDataProtocol.SELF_CLEAN_MODE: DyadProtocolCacheEntry(lambda val: DyadSelfCleanMode(val).name),
+    RoborockDyadDataProtocol.SELF_CLEAN_LEVEL: DyadProtocolCacheEntry(lambda val: DyadSelfCleanLevel(val).name),
+    RoborockDyadDataProtocol.WARM_LEVEL: DyadProtocolCacheEntry(lambda val: DyadWarmLevel(val).name),
+    RoborockDyadDataProtocol.CLEAN_MODE: DyadProtocolCacheEntry(lambda val: DyadCleanMode(val).name),
+    RoborockDyadDataProtocol.SUCTION: DyadProtocolCacheEntry(lambda val: DyadSuction(val).name),
+    RoborockDyadDataProtocol.WATER_LEVEL: DyadProtocolCacheEntry(lambda val: DyadWaterLevel(val).name),
+    RoborockDyadDataProtocol.BRUSH_SPEED: DyadProtocolCacheEntry(lambda val: DyadBrushSpeed(val).name),
+    RoborockDyadDataProtocol.POWER: DyadProtocolCacheEntry(lambda val: int(val)),
+    RoborockDyadDataProtocol.AUTO_DRY: DyadProtocolCacheEntry(lambda val: bool(val)),
+    RoborockDyadDataProtocol.MESH_LEFT: DyadProtocolCacheEntry(lambda val: int(360000 - val * 60)),
+    RoborockDyadDataProtocol.BRUSH_LEFT: DyadProtocolCacheEntry(lambda val: int(360000 - val * 60)),
+    RoborockDyadDataProtocol.ERROR: DyadProtocolCacheEntry(lambda val: DyadError(val).name),
+    RoborockDyadDataProtocol.VOLUME_SET: DyadProtocolCacheEntry(lambda val: int(val)),
+    RoborockDyadDataProtocol.STAND_LOCK_AUTO_RUN: DyadProtocolCacheEntry(lambda val: bool(val)),
+    RoborockDyadDataProtocol.AUTO_DRY_MODE: DyadProtocolCacheEntry(lambda val: bool(val)),
+    RoborockDyadDataProtocol.SILENT_DRY_DURATION: DyadProtocolCacheEntry(lambda val: int(val)),  # in minutes
+    RoborockDyadDataProtocol.SILENT_MODE: DyadProtocolCacheEntry(lambda val: bool(val)),
+    RoborockDyadDataProtocol.SILENT_MODE_START_TIME: DyadProtocolCacheEntry(
+        lambda val: time(hour=int(val / 60), minute=val % 60)
+    ),  # in minutes since 00:00
+    RoborockDyadDataProtocol.SILENT_MODE_END_TIME: DyadProtocolCacheEntry(
+        lambda val: time(hour=int(val / 60), minute=val % 60)
+    ),  # in minutes since 00:00
+    RoborockDyadDataProtocol.RECENT_RUN_TIME: DyadProtocolCacheEntry(
+        lambda val: [int(v) for v in val.split(",")]
+    ),  # minutes of cleaning in past few days.
+    RoborockDyadDataProtocol.TOTAL_RUN_TIME: DyadProtocolCacheEntry(lambda val: int(val)),
+    RoborockDyadDataProtocol.SND_STATE: DyadProtocolCacheEntry(lambda val: DyadSndState.from_dict(val)),
+    RoborockDyadDataProtocol.PRODUCT_INFO: DyadProtocolCacheEntry(lambda val: DyadProductInfo.from_dict(val)),
+}
+
+
+class RoborockClientA01(RoborockClient):
+    def __init__(self, endpoint: str, device_info: DeviceData):
+        super().__init__(endpoint, device_info)
+
+    def on_message_received(self, messages: list[RoborockMessage]) -> None:
+        for message in messages:
+            protocol = message.protocol
+            if message.payload and protocol in [
+                RoborockMessageProtocol.RPC_RESPONSE,
+                RoborockMessageProtocol.GENERAL_REQUEST,
+            ]:
+                payload = message.payload
+                try:
+                    payload = unpad(payload, AES.block_size)
+                except Exception:
+                    continue
+                payload_json = json.loads(payload.decode())
+                for data_point_number, data_point in payload_json.get("dps").items():
+                    data_point_protocol = RoborockDyadDataProtocol(int(data_point_number))
+                    if data_point_protocol in protocol_entries:
+                        # Auto convert into data struct we want.
+                        converted_response = protocol_entries[data_point_protocol].post_process_fn(data_point)
+                        queue = self._waiting_queue.get(int(data_point_number))
+                        if queue and queue.protocol == protocol:
+                            queue.resolve((converted_response, None))
+
+    async def update_values(self, dyad_data_protocols: list[RoborockDyadDataProtocol]):
+        """This should handle updating for each given protocol."""
+        raise NotImplementedError

roborock/version_a01_apis/roborock_mqtt_client_a01.py

@@ -0,0 +1,55 @@
+import asyncio
+import base64
+import json
+
+from Crypto.Cipher import AES
+from Crypto.Util.Padding import pad, unpad
+
+from roborock.cloud_api import RoborockMqttClient
+from roborock.containers import DeviceData, UserData
+from roborock.exceptions import RoborockException
+from roborock.protocol import MessageParser, Utils
+from roborock.roborock_message import RoborockDyadDataProtocol, RoborockMessage, RoborockMessageProtocol
+
+from .roborock_client_a01 import RoborockClientA01
+
+
+class RoborockMqttClientA01(RoborockMqttClient, RoborockClientA01):
+    def __init__(self, user_data: UserData, device_info: DeviceData, queue_timeout: int = 10) -> None:
+        rriot = user_data.rriot
+        if rriot is None:
+            raise RoborockException("Got no rriot data from user_data")
+        endpoint = base64.b64encode(Utils.md5(rriot.k.encode())[8:14]).decode()
+
+        RoborockMqttClient.__init__(self, user_data, device_info, queue_timeout)
+        RoborockClientA01.__init__(self, endpoint, device_info)
+
+    async def send_message(self, roborock_message: RoborockMessage):
+        await self.validate_connection()
+        response_protocol = RoborockMessageProtocol.RPC_RESPONSE
+
+        local_key = self.device_info.device.local_key
+        m = MessageParser.build(roborock_message, local_key, prefixed=False)
+        # self._logger.debug(f"id={request_id} Requesting method {method} with {params}")
+        payload = json.loads(unpad(roborock_message.payload, AES.block_size))
+        futures = []
+        if "10000" in payload["dps"]:
+            for dps in json.loads(payload["dps"]["10000"]):
+                futures.append(asyncio.ensure_future(self._async_response(dps, response_protocol)))
+        self._send_msg_raw(m)
+        responses = await asyncio.gather(*futures)
+        dps_responses = {}
+        if "10000" in payload["dps"]:
+            for i, dps in enumerate(json.loads(payload["dps"]["10000"])):
+                dps_responses[dps] = responses[i][0]
+        return dps_responses
+
+    async def update_values(self, dyad_data_protocols: list[RoborockDyadDataProtocol]):
+        payload = {"dps": {RoborockDyadDataProtocol.ID_QUERY: str([int(protocol) for protocol in dyad_data_protocols])}}
+        return await self.send_message(
+            RoborockMessage(
+                protocol=RoborockMessageProtocol.RPC_REQUEST,
+                version=b"A01",
+                payload=pad(json.dumps(payload).encode("utf-8"), AES.block_size),
+            )
+        )