python-roborock 2.35.0__tar.gz → 2.37.0__tar.gz

{python_roborock-2.35.0 → python_roborock-2.37.0}/PKG-INFO

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

{python_roborock-2.35.0 → python_roborock-2.37.0}/pyproject.toml

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

{python_roborock-2.35.0 → python_roborock-2.37.0}/roborock/api.py

@@ -3,9 +3,7 @@
 from __future__ import annotations
 
 import asyncio
-import base64
 import logging
-import secrets
 import time
 from abc import ABC, abstractmethod
 from typing import Any
@@ -37,14 +35,11 @@ class RoborockClient(ABC):
     def __init__(self, device_info: DeviceData) -> None:
         """Initialize RoborockClient."""
         self.device_info = device_info
-        self._nonce = secrets.token_bytes(16)
         self._waiting_queue: dict[int, RoborockFuture] = {}
         self._last_device_msg_in = time.monotonic()
         self._last_disconnection = time.monotonic()
         self.keep_alive = KEEPALIVE
-        self._diagnostic_data: dict[str, dict[str, Any]] = {
-            "misc_info": {"Nonce": base64.b64encode(self._nonce).decode("utf-8")}
-        }
+        self._diagnostic_data: dict[str, dict[str, Any]] = {}
         self.is_available: bool = True
 
     async def async_release(self) -> None:

{python_roborock-2.35.0 → python_roborock-2.37.0}/roborock/containers.py

@@ -725,6 +725,29 @@ class NetworkInfo(RoborockBase):
     rssi: int | None = None
 
 
+@dataclass
+class AppInitStatusLocalInfo(RoborockBase):
+    location: str
+    bom: str | None = None
+    featureset: int | None = None
+    language: str | None = None
+    logserver: str | None = None
+    wifiplan: str | None = None
+    timezone: str | None = None
+    name: str | None = None
+
+
+@dataclass
+class AppInitStatus(RoborockBase):
+    local_info: AppInitStatusLocalInfo
+    feature_info: list[int]
+    new_feature_info: int
+    new_feature_info_str: str
+    new_feature_info_2: int | None = None
+    carriage_type: int | None = None
+    dsp_version: int | None = None
+
+
 @dataclass
 class DeviceData(RoborockBase):
     device: HomeDataDevice

python_roborock-2.37.0/roborock/devices/a01_channel.py

@@ -0,0 +1,93 @@
+"""Thin wrapper around the MQTT channel for Roborock A01 devices."""
+
+import asyncio
+import logging
+from typing import Any, overload
+
+from roborock.exceptions import RoborockException
+from roborock.protocols.a01_protocol import (
+    decode_rpc_response,
+    encode_mqtt_payload,
+)
+from roborock.roborock_message import (
+    RoborockDyadDataProtocol,
+    RoborockMessage,
+    RoborockZeoProtocol,
+)
+
+from .mqtt_channel import MqttChannel
+
+_LOGGER = logging.getLogger(__name__)
+_TIMEOUT = 10.0
+
+# Both RoborockDyadDataProtocol and RoborockZeoProtocol have the same
+# value for ID_QUERY
+_ID_QUERY = int(RoborockDyadDataProtocol.ID_QUERY)
+
+
+@overload
+async def send_decoded_command(
+    mqtt_channel: MqttChannel,
+    params: dict[RoborockDyadDataProtocol, Any],
+) -> dict[RoborockDyadDataProtocol, Any]:
+    ...
+
+
+@overload
+async def send_decoded_command(
+    mqtt_channel: MqttChannel,
+    params: dict[RoborockZeoProtocol, Any],
+) -> dict[RoborockZeoProtocol, Any]:
+    ...
+
+
+async def send_decoded_command(
+    mqtt_channel: MqttChannel,
+    params: dict[RoborockDyadDataProtocol, Any] | dict[RoborockZeoProtocol, Any],
+) -> dict[RoborockDyadDataProtocol, Any] | dict[RoborockZeoProtocol, Any]:
+    """Send a command on the MQTT channel and get a decoded response."""
+    _LOGGER.debug("Sending MQTT command: %s", params)
+    roborock_message = encode_mqtt_payload(params)
+
+    # For commands that set values: send the command and do not
+    # block waiting for a response. Queries are handled below.
+    param_values = {int(k): v for k, v in params.items()}
+    if not (query_values := param_values.get(_ID_QUERY)):
+        await mqtt_channel.publish(roborock_message)
+        return {}
+
+    # Merge any results together than contain the requested data. This
+    # does not use a future since it needs to merge results across responses.
+    # This could be simplified if we can assume there is a single response.
+    finished = asyncio.Event()
+    result: dict[int, Any] = {}
+
+    def find_response(response_message: RoborockMessage) -> None:
+        """Handle incoming messages and resolve the future."""
+        try:
+            decoded = decode_rpc_response(response_message)
+        except RoborockException as ex:
+            _LOGGER.info("Failed to decode a01 message: %s: %s", response_message, ex)
+            return
+        for key, value in decoded.items():
+            if key in query_values:
+                result[key] = value
+        if len(result) != len(query_values):
+            _LOGGER.debug("Incomplete query response: %s != %s", result, query_values)
+            return
+        _LOGGER.debug("Received query response: %s", result)
+        if not finished.is_set():
+            finished.set()
+
+    unsub = await mqtt_channel.subscribe(find_response)
+
+    try:
+        await mqtt_channel.publish(roborock_message)
+        try:
+            await asyncio.wait_for(finished.wait(), timeout=_TIMEOUT)
+        except TimeoutError as ex:
+            raise RoborockException(f"Command timed out after {_TIMEOUT}s") from ex
+    finally:
+        unsub()
+
+    return result  # type: ignore[return-value]

{python_roborock-2.35.0 → python_roborock-2.37.0}/roborock/devices/b01_channel.py

@@ -3,12 +3,10 @@
 from __future__ import annotations
 
 import logging
-from typing import Any
 
 from roborock.protocols.b01_protocol import (
     CommandType,
     ParamsType,
-    decode_rpc_response,
     encode_mqtt_payload,
 )
 
@@ -22,9 +20,8 @@ async def send_decoded_command(
     dps: int,
     command: CommandType,
     params: ParamsType,
-) -> dict[int, Any]:
+) -> None:
     """Send a command on the MQTT channel and get a decoded response."""
     _LOGGER.debug("Sending MQTT command: %s", params)
     roborock_message = encode_mqtt_payload(dps, command, params)
-    response = await mqtt_channel.send_message(roborock_message)
-    return decode_rpc_response(response)  # type: ignore[return-value]
+    await mqtt_channel.publish(roborock_message)

{python_roborock-2.35.0 → python_roborock-2.37.0}/roborock/devices/local_channel.py

@@ -4,7 +4,6 @@ import asyncio
 import logging
 from collections.abc import Callable
 from dataclasses import dataclass
-from json import JSONDecodeError
 
 from roborock.exceptions import RoborockConnectionException, RoborockException
 from roborock.protocol import Decoder, Encoder, create_local_decoder, create_local_encoder
@@ -46,11 +45,8 @@ class LocalChannel(Channel):
         self._subscribers: list[Callable[[RoborockMessage], None]] = []
         self._is_connected = False
 
-        # RPC support
-        self._waiting_queue: dict[int, asyncio.Future[RoborockMessage]] = {}
         self._decoder: Decoder = create_local_decoder(local_key)
         self._encoder: Encoder = create_local_encoder(local_key)
-        self._queue_lock = asyncio.Lock()
 
     @property
     def is_connected(self) -> bool:
@@ -87,7 +83,6 @@ class LocalChannel(Channel):
             return
         for message in messages:
             _LOGGER.debug("Received message: %s", message)
-            asyncio.create_task(self._resolve_future_with_lock(message))
             for callback in self._subscribers:
                 try:
                     callback(message)
@@ -109,48 +104,24 @@ class LocalChannel(Channel):
 
         return unsubscribe
 
-    async def _resolve_future_with_lock(self, message: RoborockMessage) -> None:
-        """Resolve waiting future with proper locking."""
-        if (request_id := message.get_request_id()) is None:
-            _LOGGER.debug("Received message with no request_id")
-            return
-        async with self._queue_lock:
-            if (future := self._waiting_queue.pop(request_id, None)) is not None:
-                future.set_result(message)
-            else:
-                _LOGGER.debug("Received message with no waiting handler: request_id=%s", request_id)
-
-    async def send_message(self, message: RoborockMessage, timeout: float = 10.0) -> RoborockMessage:
-        """Send a command message and wait for the response message."""
+    async def publish(self, message: RoborockMessage) -> None:
+        """Send a command message.
+
+        The caller is responsible for associating the message with its response.
+        """
         if not self._transport or not self._is_connected:
             raise RoborockConnectionException("Not connected to device")
 
-        try:
-            if (request_id := message.get_request_id()) is None:
-                raise RoborockException("Message must have a request_id for RPC calls")
-        except (ValueError, JSONDecodeError) as err:
-            _LOGGER.exception("Error getting request_id from message: %s", err)
-            raise RoborockException(f"Invalid message format, Message must have a request_id: {err}") from err
-
-        future: asyncio.Future[RoborockMessage] = asyncio.Future()
-        async with self._queue_lock:
-            if request_id in self._waiting_queue:
-                raise RoborockException(f"Request ID {request_id} already pending, cannot send command")
-            self._waiting_queue[request_id] = future
-
         try:
             encoded_msg = self._encoder(message)
+        except Exception as err:
+            _LOGGER.exception("Error encoding MQTT message: %s", err)
+            raise RoborockException(f"Failed to encode MQTT message: {err}") from err
+        try:
             self._transport.write(encoded_msg)
-            return await asyncio.wait_for(future, timeout=timeout)
-        except asyncio.TimeoutError as ex:
-            async with self._queue_lock:
-                self._waiting_queue.pop(request_id, None)
-            raise RoborockException(f"Command timed out after {timeout}s") from ex
-        except Exception:
+        except Exception as err:
             logging.exception("Uncaught error sending command")
-            async with self._queue_lock:
-                self._waiting_queue.pop(request_id, None)
-            raise
+            raise RoborockException(f"Failed to send message: {message}") from err
 
 
 # This module provides a factory function to create LocalChannel instances.

python_roborock-2.37.0/roborock/devices/mqtt_channel.py

@@ -0,0 +1,95 @@
+"""Modules for communicating with specific Roborock devices over MQTT."""
+
+import logging
+from collections.abc import Callable
+
+from roborock.containers import HomeDataDevice, RRiot, UserData
+from roborock.exceptions import RoborockException
+from roborock.mqtt.session import MqttParams, MqttSession, MqttSessionException
+from roborock.protocol import create_mqtt_decoder, create_mqtt_encoder
+from roborock.roborock_message import RoborockMessage
+
+from .channel import Channel
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class MqttChannel(Channel):
+    """Simple RPC-style channel for communicating with a device over MQTT.
+
+    Handles request/response correlation and timeouts, but leaves message
+    format most parsing to higher-level components.
+    """
+
+    def __init__(self, mqtt_session: MqttSession, duid: str, local_key: str, rriot: RRiot, mqtt_params: MqttParams):
+        self._mqtt_session = mqtt_session
+        self._duid = duid
+        self._local_key = local_key
+        self._rriot = rriot
+        self._mqtt_params = mqtt_params
+
+        self._decoder = create_mqtt_decoder(local_key)
+        self._encoder = create_mqtt_encoder(local_key)
+
+    @property
+    def is_connected(self) -> bool:
+        """Return true if the channel is connected.
+
+        This passes through the underlying MQTT session's connected state.
+        """
+        return self._mqtt_session.connected
+
+    @property
+    def _publish_topic(self) -> str:
+        """Topic to send commands to the device."""
+        return f"rr/m/i/{self._rriot.u}/{self._mqtt_params.username}/{self._duid}"
+
+    @property
+    def _subscribe_topic(self) -> str:
+        """Topic to receive responses from the device."""
+        return f"rr/m/o/{self._rriot.u}/{self._mqtt_params.username}/{self._duid}"
+
+    async def subscribe(self, callback: Callable[[RoborockMessage], None]) -> Callable[[], None]:
+        """Subscribe to the device's response topic.
+
+        The callback will be called with the message payload when a message is received.
+
+        Returns a callable that can be used to unsubscribe from the topic.
+        """
+
+        def message_handler(payload: bytes) -> None:
+            if not (messages := self._decoder(payload)):
+                _LOGGER.warning("Failed to decode MQTT message: %s", payload)
+                return
+            for message in messages:
+                _LOGGER.debug("Received message: %s", message)
+                try:
+                    callback(message)
+                except Exception as e:
+                    _LOGGER.exception("Uncaught error in message handler callback: %s", e)
+
+        return await self._mqtt_session.subscribe(self._subscribe_topic, message_handler)
+
+    async def publish(self, message: RoborockMessage) -> None:
+        """Publish a command message.
+
+        The caller is responsible for handling any responses and associating them
+        with the incoming request.
+        """
+        try:
+            encoded_msg = self._encoder(message)
+        except Exception as e:
+            _LOGGER.exception("Error encoding MQTT message: %s", e)
+            raise RoborockException(f"Failed to encode MQTT message: {e}") from e
+        try:
+            return await self._mqtt_session.publish(self._publish_topic, encoded_msg)
+        except MqttSessionException as e:
+            _LOGGER.exception("Error publishing MQTT message: %s", e)
+            raise RoborockException(f"Failed to publish MQTT message: {e}") from e
+
+
+def create_mqtt_channel(
+    user_data: UserData, mqtt_params: MqttParams, mqtt_session: MqttSession, device: HomeDataDevice
+) -> MqttChannel:
+    """Create a V1Channel for the given device."""
+    return MqttChannel(mqtt_session, device.duid, device.local_key, user_data.rriot, mqtt_params)

{python_roborock-2.35.0 → python_roborock-2.37.0}/roborock/devices/traits/b01/props.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import logging
-from typing import Any
 
 from roborock import RoborockB01Methods
 from roborock.roborock_message import RoborockB01Props
@@ -26,6 +25,6 @@ class B01PropsApi(Trait):
         """Initialize the B01Props API."""
         self._channel = channel
 
-    async def query_values(self, props: list[RoborockB01Props]) -> dict[int, Any]:
+    async def query_values(self, props: list[RoborockB01Props]) -> None:
         """Query the device for the values of the given Dyad protocols."""
-        return await send_decoded_command(self._channel, dps=10000, command=RoborockB01Methods.GET_PROP, params=props)
+        await send_decoded_command(self._channel, dps=10000, command=RoborockB01Methods.GET_PROP, params=props)

{python_roborock-2.35.0 → python_roborock-2.37.0}/roborock/devices/v1_channel.py

@@ -79,7 +79,7 @@ class V1Channel(Channel):
     @property
     def is_mqtt_connected(self) -> bool:
         """Return whether MQTT connection is available."""
-        return self._mqtt_unsub is not None
+        return self._mqtt_unsub is not None and self._mqtt_channel.is_connected
 
     @property
     def rpc_channel(self) -> V1RpcChannel:

{python_roborock-2.35.0 → python_roborock-2.37.0}/roborock/devices/v1_rpc_channel.py

@@ -6,25 +6,27 @@ a simple interface for sending commands and receiving responses over both MQTT
 and local connections, preferring local when available.
 """
 
+import asyncio
 import logging
 from collections.abc import Callable
 from typing import Any, Protocol, TypeVar, overload
 
 from roborock.containers import RoborockBase
+from roborock.exceptions import RoborockException
 from roborock.protocols.v1_protocol import (
     CommandType,
     ParamsType,
+    RequestMessage,
     SecurityData,
-    create_mqtt_payload_encoder,
     decode_rpc_response,
-    encode_local_payload,
 )
-from roborock.roborock_message import RoborockMessage
+from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
 
 from .local_channel import LocalChannel
 from .mqtt_channel import MqttChannel
 
 _LOGGER = logging.getLogger(__name__)
+_TIMEOUT = 10.0
 
 
 _T = TypeVar("_T", bound=RoborockBase)
@@ -116,7 +118,7 @@ class PayloadEncodedV1RpcChannel(BaseV1RpcChannel):
         self,
         name: str,
         channel: MqttChannel | LocalChannel,
-        payload_encoder: Callable[[CommandType, ParamsType], RoborockMessage],
+        payload_encoder: Callable[[RequestMessage], RoborockMessage],
     ) -> None:
         """Initialize the channel with a raw channel and an encoder function."""
         self._name = name
@@ -131,18 +133,44 @@ class PayloadEncodedV1RpcChannel(BaseV1RpcChannel):
     ) -> Any:
         """Send a command and return a parsed response RoborockBase type."""
         _LOGGER.debug("Sending command (%s): %s, params=%s", self._name, method, params)
-        message = self._payload_encoder(method, params)
-        response = await self._channel.send_message(message)
-        return decode_rpc_response(response)
+        request_message = RequestMessage(method, params=params)
+        message = self._payload_encoder(request_message)
+
+        future: asyncio.Future[dict[str, Any]] = asyncio.Future()
+
+        def find_response(response_message: RoborockMessage) -> None:
+            try:
+                decoded = decode_rpc_response(response_message)
+            except RoborockException:
+                return
+            if decoded.request_id == request_message.request_id:
+                future.set_result(decoded.data)
+
+        unsub = await self._channel.subscribe(find_response)
+        try:
+            await self._channel.publish(message)
+            return await asyncio.wait_for(future, timeout=_TIMEOUT)
+        except TimeoutError as ex:
+            future.cancel()
+            raise RoborockException(f"Command timed out after {_TIMEOUT}s") from ex
+        finally:
+            unsub()
 
 
 def create_mqtt_rpc_channel(mqtt_channel: MqttChannel, security_data: SecurityData) -> V1RpcChannel:
     """Create a V1 RPC channel using an MQTT channel."""
-    payload_encoder = create_mqtt_payload_encoder(security_data)
-    return PayloadEncodedV1RpcChannel("mqtt", mqtt_channel, payload_encoder)
+    return PayloadEncodedV1RpcChannel(
+        "mqtt",
+        mqtt_channel,
+        lambda x: x.encode_message(RoborockMessageProtocol.RPC_REQUEST, security_data=security_data),
+    )
 
 
 def create_combined_rpc_channel(local_channel: LocalChannel, mqtt_rpc_channel: V1RpcChannel) -> V1RpcChannel:
     """Create a V1 RPC channel that combines local and MQTT channels."""
-    local_rpc_channel = PayloadEncodedV1RpcChannel("local", local_channel, encode_local_payload)
+    local_rpc_channel = PayloadEncodedV1RpcChannel(
+        "local",
+        local_channel,
+        lambda x: x.encode_message(RoborockMessageProtocol.GENERAL_REQUEST),
+    )
     return CombinedV1RpcChannel(local_channel, local_rpc_channel, mqtt_rpc_channel)

{python_roborock-2.35.0 → python_roborock-2.37.0}/roborock/protocol.py

@@ -147,6 +147,26 @@ class Utils:
             return unpad(decipher.decrypt(ciphertext), AES.block_size)
         return ciphertext
 
+    @staticmethod
+    def encrypt_cbc(plaintext: bytes, token: bytes) -> bytes:
+        """Encrypt plaintext with a given token using cbc mode.
+
+        This is currently used for testing purposes only.
+
+        :param bytes plaintext: Plaintext (json) to encrypt
+        :param bytes token: Token to use
+        :return: Encrypted bytes
+        """
+        if not isinstance(plaintext, bytes):
+            raise TypeError("plaintext requires bytes")
+        Utils.verify_token(token)
+        iv = bytes(AES.block_size)
+        cipher = AES.new(token, AES.MODE_CBC, iv)
+        if plaintext:
+            plaintext = pad(plaintext, AES.block_size)
+            return cipher.encrypt(plaintext)
+        return plaintext
+
     @staticmethod
     def decrypt_cbc(ciphertext: bytes, token: bytes) -> bytes:
         """Decrypt ciphertext with a given token using cbc mode.

{python_roborock-2.35.0 → python_roborock-2.37.0}/roborock/protocols/v1_protocol.py

@@ -7,6 +7,7 @@ import json
 import logging
 import math
 import secrets
+import struct
 import time
 from collections.abc import Callable
 from dataclasses import dataclass, field
@@ -24,8 +25,6 @@ _LOGGER = logging.getLogger(__name__)
 __all__ = [
     "SecurityData",
     "create_security_data",
-    "create_mqtt_payload_encoder",
-    "encode_local_payload",
     "decode_rpc_response",
 ]
 
@@ -44,6 +43,10 @@ class SecurityData:
         """Convert security data to a dictionary for sending in the payload."""
         return {"security": {"endpoint": self.endpoint, "nonce": self.nonce.hex().lower()}}
 
+    def to_diagnostic_data(self) -> dict[str, Any]:
+        """Convert security data to a dictionary for debugging purposes."""
+        return {"nonce": self.nonce.hex().lower()}
+
 
 def create_security_data(rriot: RRiot) -> SecurityData:
     """Create a SecurityData instance for the given endpoint and nonce."""
@@ -61,7 +64,19 @@ class RequestMessage:
     timestamp: int = field(default_factory=lambda: math.floor(time.time()))
     request_id: int = field(default_factory=lambda: get_next_int(10000, 32767))
 
-    def as_payload(self, security_data: SecurityData | None) -> bytes:
+    def encode_message(
+        self,
+        protocol: RoborockMessageProtocol,
+        security_data: SecurityData | None = None,
+    ) -> RoborockMessage:
+        """Convert the request message to a RoborockMessage."""
+        return RoborockMessage(
+            timestamp=self.timestamp,
+            protocol=protocol,
+            payload=self._as_payload(security_data=security_data),
+        )
+
+    def _as_payload(self, security_data: SecurityData | None) -> bytes:
         """Convert the request arguments to a dictionary."""
         inner = {
             "id": self.request_id,
@@ -80,36 +95,18 @@ class RequestMessage:
         )
 
 
-def create_mqtt_payload_encoder(security_data: SecurityData) -> Callable[[CommandType, ParamsType], RoborockMessage]:
-    """Create a payload encoder for V1 commands over MQTT."""
-
-    def _get_payload(method: CommandType, params: ParamsType) -> RoborockMessage:
-        """Build the payload for a V1 command."""
-        request = RequestMessage(method=method, params=params)
-        payload = request.as_payload(security_data)  # always secure
-        return RoborockMessage(
-            timestamp=request.timestamp,
-            protocol=RoborockMessageProtocol.RPC_REQUEST,
-            payload=payload,
-        )
-
-    return _get_payload
-
+@dataclass(kw_only=True, frozen=True)
+class ResponseMessage:
+    """Data structure for v1 RoborockMessage responses."""
 
-def encode_local_payload(method: CommandType, params: ParamsType) -> RoborockMessage:
-    """Encode payload for V1 commands over local connection."""
+    request_id: int | None
+    """The request ID of the response."""
 
-    request = RequestMessage(method=method, params=params)
-    payload = request.as_payload(security_data=None)
+    data: dict[str, Any]
+    """The data of the response."""
 
-    return RoborockMessage(
-        timestamp=request.timestamp,
-        protocol=RoborockMessageProtocol.GENERAL_REQUEST,
-        payload=payload,
-    )
 
-
-def decode_rpc_response(message: RoborockMessage) -> dict[str, Any]:
+def decode_rpc_response(message: RoborockMessage) -> ResponseMessage:
     """Decode a V1 RPC_RESPONSE message."""
     if not message.payload:
         raise RoborockException("Invalid V1 message format: missing payload")
@@ -123,14 +120,19 @@ def decode_rpc_response(message: RoborockMessage) -> dict[str, Any]:
     if not isinstance(datapoints, dict):
         raise RoborockException(f"Invalid V1 message format: 'dps' should be a dictionary for {message.payload!r}")
 
-    if not (data_point := datapoints.get("102")):
-        raise RoborockException("Invalid V1 message format: missing '102' data point")
+    if not (data_point := datapoints.get(str(RoborockMessageProtocol.RPC_RESPONSE))):
+        raise RoborockException(
+            f"Invalid V1 message format: missing '{RoborockMessageProtocol.RPC_RESPONSE}' data point"
+        )
 
     try:
         data_point_response = json.loads(data_point)
     except (json.JSONDecodeError, TypeError) as e:
-        raise RoborockException(f"Invalid V1 message data point '102': {e} for {message.payload!r}") from e
+        raise RoborockException(
+            f"Invalid V1 message data point '{RoborockMessageProtocol.RPC_RESPONSE}': {e} for {message.payload!r}"
+        ) from e
 
+    request_id: int | None = data_point_response.get("id")
     if error := data_point_response.get("error"):
         raise RoborockException(f"Error in message: {error}")
 
@@ -141,4 +143,38 @@ def decode_rpc_response(message: RoborockMessage) -> dict[str, Any]:
         result = result[0]
     if not isinstance(result, dict):
         raise RoborockException(f"Invalid V1 message format: 'result' should be a dictionary for {message.payload!r}")
-    return result
+    return ResponseMessage(request_id=request_id, data=result)
+
+
+@dataclass
+class MapResponse:
+    """Data structure for the V1 Map response."""
+
+    request_id: int
+    """The request ID of the map response."""
+
+    data: bytes
+    """The map data, decrypted and decompressed."""
+
+
+def create_map_response_decoder(security_data: SecurityData) -> Callable[[RoborockMessage], MapResponse]:
+    """Create a decoder for V1 map response messages."""
+
+    def _decode_map_response(message: RoborockMessage) -> MapResponse:
+        """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}"
+            )
+        try:
+            decrypted = Utils.decrypt_cbc(body, security_data.nonce)
+        except ValueError as err:
+            raise RoborockException("Failed to decode map message payload") from err
+        decompressed = Utils.decompress(decrypted)
+        return MapResponse(request_id=request_id, data=decompressed)
+
+    return _decode_map_response

{python_roborock-2.35.0 → python_roborock-2.37.0}/roborock/roborock_message.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import json
 import math
 import time
 from dataclasses import dataclass, field
@@ -247,12 +246,3 @@ class RoborockMessage:
     version: bytes = b"1.0"
     random: int = field(default_factory=lambda: get_next_int(10000, 99999))
     timestamp: int = field(default_factory=lambda: math.floor(time.time()))
-
-    def get_request_id(self) -> int | None:
-        if self.payload:
-            payload = json.loads(self.payload.decode())
-            for data_point_number, data_point in payload.get("dps").items():
-                if data_point_number in ["101", "102"]:
-                    data_point_response = json.loads(data_point)
-                    return data_point_response.get("id")
-        return None

{python_roborock-2.35.0 → python_roborock-2.37.0}/roborock/version_1_apis/roborock_client_v1.py

@@ -1,13 +1,13 @@
 import asyncio
 import dataclasses
 import json
-import struct
 import time
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Coroutine
 from typing import Any, TypeVar, final
 
 from roborock import (
+    AppInitStatus,
     DeviceProp,
     DockSummary,
     RoborockCommand,
@@ -45,7 +45,7 @@ from roborock.containers import (
     ValleyElectricityTimer,
     WashTowelMode,
 )
-from roborock.protocol import Utils
+from roborock.protocols.v1_protocol import MapResponse, SecurityData, create_map_response_decoder
 from roborock.roborock_message import (
     ROBOROCK_DATA_CONSUMABLE_PROTOCOL,
     ROBOROCK_DATA_STATUS_PROTOCOL,
@@ -150,10 +150,15 @@ class RoborockClientV1(RoborockClient, ABC):
     """Roborock client base class for version 1 devices."""
 
     _listeners: dict[str, ListenerModel] = {}
+    _map_response_decoder: Callable[[RoborockMessage], MapResponse] | None = None
 
-    def __init__(self, device_info: DeviceData, endpoint: str):
+    def __init__(self, device_info: DeviceData, security_data: SecurityData | None) -> None:
         """Initializes the Roborock client."""
         super().__init__(device_info)
+        if security_data is not None:
+            self._diagnostic_data.update({"misc_info": security_data.to_diagnostic_data()})
+            self._map_response_decoder = create_map_response_decoder(security_data)
+
         self._status_type: type[Status] = ModelStatus.get(device_info.model, S7MaxVStatus)
         self.cache: dict[CacheableAttribute, AttributeCache] = {
             cacheable_attribute: AttributeCache(attr, self._send_command)
@@ -162,7 +167,6 @@ class RoborockClientV1(RoborockClient, ABC):
         if device_info.device.duid not in self._listeners:
             self._listeners[device_info.device.duid] = ListenerModel({}, self.cache)
         self.listener_model = self._listeners[device_info.device.duid]
-        self._endpoint = endpoint
 
     async def async_release(self) -> None:
         await super().async_release()
@@ -339,6 +343,10 @@ class RoborockClientV1(RoborockClient, ABC):
         """Load the map into the vacuum's memory."""
         await self.send_command(RoborockCommand.LOAD_MULTI_MAP, [map_flag])
 
+    async def get_app_init_status(self) -> AppInitStatus:
+        """Gets the app init status (needed for determining vacuum capabilities)."""
+        return await self.send_command(RoborockCommand.APP_GET_INIT_STATUS, return_type=AppInitStatus)
+
     @abstractmethod
     async def _send_command(
         self,
@@ -429,21 +437,15 @@ class RoborockClientV1(RoborockClient, ABC):
                             dps = {data_point_number: data_point}
                             self._logger.debug(f"Got unknown data point {dps}")
                 elif data.payload and protocol == RoborockMessageProtocol.MAP_RESPONSE:
-                    payload = data.payload[0:24]
-                    [endpoint, _, request_id, _] = struct.unpack("<8s8sH6s", payload)
-                    if endpoint.decode().startswith(self._endpoint):
-                        try:
-                            decrypted = Utils.decrypt_cbc(data.payload[24:], self._nonce)
-                        except ValueError as err:
-                            raise RoborockException(f"Failed to decode {data.payload!r} for {data.protocol}") from err
-                        decompressed = Utils.decompress(decrypted)
-                        queue = self._waiting_queue.get(request_id)
+                    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:
-                            if isinstance(decompressed, list):
-                                decompressed = decompressed[0]
-                            queue.set_result(decompressed)
+                            queue.set_result(map_response.data)
                         else:
-                            self._logger.debug("Received response for unknown request id %s", request_id)
+                            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.35.0 → python_roborock-2.37.0}/roborock/version_1_apis/roborock_local_client_v1.py

@@ -10,7 +10,7 @@ from .. import CommandVacuumError, DeviceData, RoborockCommand
 from ..api import RoborockClient
 from ..exceptions import RoborockConnectionException, RoborockException, VacuumError
 from ..protocol import Decoder, Encoder, create_local_decoder, create_local_encoder
-from ..protocols.v1_protocol import encode_local_payload
+from ..protocols.v1_protocol import RequestMessage
 from ..roborock_message import RoborockMessage, RoborockMessageProtocol
 from ..util import RoborockLoggerAdapter
 from .roborock_client_v1 import CLOUD_REQUIRED, RoborockClientV1
@@ -60,7 +60,7 @@ class RoborockLocalClientV1(RoborockClientV1, RoborockClient):
         self.transport: Transport | None = None
         self._mutex = Lock()
         self.keep_alive_task: TimerHandle | None = None
-        RoborockClientV1.__init__(self, device_data, "abc")
+        RoborockClientV1.__init__(self, device_data, security_data=None)
         RoborockClient.__init__(self, device_data)
         self._local_protocol = _LocalProtocol(self._data_received, self._connection_lost)
         self._encoder: Encoder = create_local_encoder(device_data.device.local_key)
@@ -123,12 +123,20 @@ class RoborockLocalClientV1(RoborockClientV1, RoborockClient):
 
     async def hello(self):
         try:
-            return await self._send_message(_HELLO_REQUEST_MESSAGE)
+            return await self._send_message(
+                roborock_message=_HELLO_REQUEST_MESSAGE,
+                request_id=_HELLO_REQUEST_MESSAGE.seq,
+                response_protocol=RoborockMessageProtocol.HELLO_RESPONSE,
+            )
         except Exception as e:
             self._logger.error(e)
 
     async def ping(self) -> None:
-        await self._send_message(_PING_REQUEST_MESSAGE)
+        await self._send_message(
+            roborock_message=_PING_REQUEST_MESSAGE,
+            request_id=_PING_REQUEST_MESSAGE.seq,
+            response_protocol=RoborockMessageProtocol.PING_RESPONSE,
+        )
 
     def _send_msg_raw(self, data: bytes):
         try:
@@ -145,27 +153,26 @@ class RoborockLocalClientV1(RoborockClientV1, RoborockClient):
     ):
         if method in CLOUD_REQUIRED:
             raise RoborockException(f"Method {method} is not supported over local connection")
-
-        roborock_message = encode_local_payload(method, params)
-        self._logger.debug("Building message id %s for method %s", roborock_message.get_request_id(), method)
-        return await self._send_message(roborock_message, method, params)
+        request_message = RequestMessage(method=method, params=params)
+        roborock_message = request_message.encode_message(RoborockMessageProtocol.GENERAL_REQUEST)
+        self._logger.debug("Building message id %s for method %s", request_message.request_id, method)
+        return await self._send_message(
+            roborock_message,
+            request_id=request_message.request_id,
+            response_protocol=RoborockMessageProtocol.GENERAL_REQUEST,
+            method=method,
+            params=params,
+        )
 
     async def _send_message(
         self,
         roborock_message: RoborockMessage,
+        request_id: int,
+        response_protocol: int,
         method: str | None = None,
         params: list | dict | int | None = None,
     ) -> RoborockMessage:
         await self.validate_connection()
-        request_id: int | None
-        if not method or not method.startswith("get"):
-            request_id = roborock_message.seq
-            response_protocol = request_id + 1
-        else:
-            request_id = roborock_message.get_request_id()
-            response_protocol = RoborockMessageProtocol.GENERAL_REQUEST
-        if request_id is None:
-            raise RoborockException(f"Failed build message {roborock_message}")
         msg = self._encoder(roborock_message)
         if method:
             self._logger.debug(f"id={request_id} Requesting method {method} with {params}")

{python_roborock-2.35.0 → python_roborock-2.37.0}/roborock/version_1_apis/roborock_mqtt_client_v1.py

@@ -1,4 +1,3 @@
-import base64
 import logging
 
 from vacuum_map_parser_base.config.color import ColorsPalette
@@ -10,8 +9,7 @@ from roborock.cloud_api import RoborockMqttClient
 
 from ..containers import DeviceData, UserData
 from ..exceptions import CommandVacuumError, RoborockException, VacuumError
-from ..protocol import Utils
-from ..protocols.v1_protocol import SecurityData, create_mqtt_payload_encoder
+from ..protocols.v1_protocol import RequestMessage, create_security_data
 from ..roborock_message import (
     RoborockMessageProtocol,
 )
@@ -30,15 +28,12 @@ class RoborockMqttClientV1(RoborockMqttClient, RoborockClientV1):
         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)
-        RoborockClientV1.__init__(self, device_info, endpoint)
+        security_data = create_security_data(rriot)
+        RoborockClientV1.__init__(self, device_info, security_data=security_data)
         self.queue_timeout = queue_timeout
         self._logger = RoborockLoggerAdapter(device_info.device.name, _LOGGER)
-        self._payload_encoder = create_mqtt_payload_encoder(
-            SecurityData(endpoint=self._endpoint, nonce=self._nonce),
-        )
+        self._security_data = security_data
 
     async def _send_command(
         self,
@@ -49,13 +44,15 @@ class RoborockMqttClientV1(RoborockMqttClient, RoborockClientV1):
             # When we have more custom commands do something more complicated here
             return await self._get_calibration_points()
 
-        roborock_message = self._payload_encoder(method, params)
-        self._logger.debug("Building message id %s for method %s", roborock_message.get_request_id, method)
+        request_message = RequestMessage(method=method, params=params)
+        roborock_message = request_message.encode_message(
+            RoborockMessageProtocol.RPC_REQUEST,
+            security_data=self._security_data,
+        )
+        self._logger.debug("Building message id %s for method %s", request_message.request_id, method)
 
         await self.validate_connection()
-        request_id = roborock_message.get_request_id()
-        if request_id is None:
-            raise RoborockException(f"Failed build message {roborock_message}")
+        request_id = request_message.request_id
         response_protocol = (
             RoborockMessageProtocol.MAP_RESPONSE if method in COMMANDS_SECURED else RoborockMessageProtocol.RPC_RESPONSE
         )

python_roborock-2.35.0/roborock/devices/a01_channel.py

@@ -1,43 +0,0 @@
-"""Thin wrapper around the MQTT channel for Roborock A01 devices."""
-
-from __future__ import annotations
-
-import logging
-from typing import Any, overload
-
-from roborock.protocols.a01_protocol import (
-    decode_rpc_response,
-    encode_mqtt_payload,
-)
-from roborock.roborock_message import RoborockDyadDataProtocol, RoborockZeoProtocol
-
-from .mqtt_channel import MqttChannel
-
-_LOGGER = logging.getLogger(__name__)
-
-
-@overload
-async def send_decoded_command(
-    mqtt_channel: MqttChannel,
-    params: dict[RoborockDyadDataProtocol, Any],
-) -> dict[RoborockDyadDataProtocol, Any]:
-    ...
-
-
-@overload
-async def send_decoded_command(
-    mqtt_channel: MqttChannel,
-    params: dict[RoborockZeoProtocol, Any],
-) -> dict[RoborockZeoProtocol, Any]:
-    ...
-
-
-async def send_decoded_command(
-    mqtt_channel: MqttChannel,
-    params: dict[RoborockDyadDataProtocol, Any] | dict[RoborockZeoProtocol, Any],
-) -> dict[RoborockDyadDataProtocol, Any] | dict[RoborockZeoProtocol, Any]:
-    """Send a command on the MQTT channel and get a decoded response."""
-    _LOGGER.debug("Sending MQTT command: %s", params)
-    roborock_message = encode_mqtt_payload(params)
-    response = await mqtt_channel.send_message(roborock_message)
-    return decode_rpc_response(response)  # type: ignore[return-value]

python_roborock-2.35.0/roborock/devices/mqtt_channel.py

@@ -1,137 +0,0 @@
-"""Modules for communicating with specific Roborock devices over MQTT."""
-
-import asyncio
-import logging
-from collections.abc import Callable
-from json import JSONDecodeError
-
-from roborock.containers import HomeDataDevice, RRiot, UserData
-from roborock.exceptions import RoborockException
-from roborock.mqtt.session import MqttParams, MqttSession
-from roborock.protocol import create_mqtt_decoder, create_mqtt_encoder
-from roborock.roborock_message import RoborockMessage
-
-from .channel import Channel
-
-_LOGGER = logging.getLogger(__name__)
-
-
-class MqttChannel(Channel):
-    """Simple RPC-style channel for communicating with a device over MQTT.
-
-    Handles request/response correlation and timeouts, but leaves message
-    format most parsing to higher-level components.
-    """
-
-    def __init__(self, mqtt_session: MqttSession, duid: str, local_key: str, rriot: RRiot, mqtt_params: MqttParams):
-        self._mqtt_session = mqtt_session
-        self._duid = duid
-        self._local_key = local_key
-        self._rriot = rriot
-        self._mqtt_params = mqtt_params
-
-        # RPC support
-        self._waiting_queue: dict[int, asyncio.Future[RoborockMessage]] = {}
-        self._decoder = create_mqtt_decoder(local_key)
-        self._encoder = create_mqtt_encoder(local_key)
-        self._queue_lock = asyncio.Lock()
-        self._mqtt_unsub: Callable[[], None] | None = None
-
-    @property
-    def is_connected(self) -> bool:
-        """Return true if the channel is connected."""
-        return (self._mqtt_unsub is not None) and self._mqtt_session.connected
-
-    @property
-    def _publish_topic(self) -> str:
-        """Topic to send commands to the device."""
-        return f"rr/m/i/{self._rriot.u}/{self._mqtt_params.username}/{self._duid}"
-
-    @property
-    def _subscribe_topic(self) -> str:
-        """Topic to receive responses from the device."""
-        return f"rr/m/o/{self._rriot.u}/{self._mqtt_params.username}/{self._duid}"
-
-    async def subscribe(self, callback: Callable[[RoborockMessage], None]) -> Callable[[], None]:
-        """Subscribe to the device's response topic.
-
-        The callback will be called with the message payload when a message is received.
-
-        All messages received will be processed through the provided callback, even
-        those sent in response to the `send_command` command.
-
-        Returns a callable that can be used to unsubscribe from the topic.
-        """
-
-        def message_handler(payload: bytes) -> None:
-            if not (messages := self._decoder(payload)):
-                _LOGGER.warning("Failed to decode MQTT message: %s", payload)
-                return
-            for message in messages:
-                _LOGGER.debug("Received message: %s", message)
-                asyncio.create_task(self._resolve_future_with_lock(message))
-                try:
-                    callback(message)
-                except Exception as e:
-                    _LOGGER.exception("Uncaught error in message handler callback: %s", e)
-
-        self._mqtt_unsub = await self._mqtt_session.subscribe(self._subscribe_topic, message_handler)
-
-        def unsub_wrapper() -> None:
-            if self._mqtt_unsub is not None:
-                self._mqtt_unsub()
-                self._mqtt_unsub = None
-
-        return unsub_wrapper
-
-    async def _resolve_future_with_lock(self, message: RoborockMessage) -> None:
-        """Resolve waiting future with proper locking."""
-        if (request_id := message.get_request_id()) is None:
-            _LOGGER.debug("Received message with no request_id")
-            return
-        async with self._queue_lock:
-            if (future := self._waiting_queue.pop(request_id, None)) is not None:
-                future.set_result(message)
-            else:
-                _LOGGER.debug("Received message with no waiting handler: request_id=%s", request_id)
-
-    async def send_message(self, message: RoborockMessage, timeout: float = 10.0) -> RoborockMessage:
-        """Send a command message and wait for the response message.
-
-        Returns the raw response message - caller is responsible for parsing.
-        """
-        try:
-            if (request_id := message.get_request_id()) is None:
-                raise RoborockException("Message must have a request_id for RPC calls")
-        except (ValueError, JSONDecodeError) as err:
-            _LOGGER.exception("Error getting request_id from message: %s", err)
-            raise RoborockException(f"Invalid message format, Message must have a request_id: {err}") from err
-
-        future: asyncio.Future[RoborockMessage] = asyncio.Future()
-        async with self._queue_lock:
-            if request_id in self._waiting_queue:
-                raise RoborockException(f"Request ID {request_id} already pending, cannot send command")
-            self._waiting_queue[request_id] = future
-
-        try:
-            encoded_msg = self._encoder(message)
-            await self._mqtt_session.publish(self._publish_topic, encoded_msg)
-
-            return await asyncio.wait_for(future, timeout=timeout)
-
-        except asyncio.TimeoutError as ex:
-            async with self._queue_lock:
-                self._waiting_queue.pop(request_id, None)
-            raise RoborockException(f"Command timed out after {timeout}s") from ex
-        except Exception:
-            logging.exception("Uncaught error sending command")
-            async with self._queue_lock:
-                self._waiting_queue.pop(request_id, None)
-            raise
-
-
-def create_mqtt_channel(
-    user_data: UserData, mqtt_params: MqttParams, mqtt_session: MqttSession, device: HomeDataDevice
-) -> MqttChannel:
-    """Create a V1Channel for the given device."""
-    return MqttChannel(mqtt_session, device.duid, device.local_key, user_data.rriot, mqtt_params)