roborock-cli 0.1.1__py3-none-any.whl

roborock_cli/_vendor/roborock/protocols/a01_protocol.py

@@ -0,0 +1,74 @@
+"""Roborock A01 Protocol encoding and decoding."""
+
+import json
+import logging
+from collections.abc import Callable
+from typing import Any
+
+from Crypto.Cipher import AES
+from Crypto.Util.Padding import pad, unpad
+
+from roborock_cli._vendor.roborock.exceptions import RoborockException
+from roborock_cli._vendor.roborock.roborock_message import (
+    RoborockDyadDataProtocol,
+    RoborockMessage,
+    RoborockMessageProtocol,
+    RoborockZeoProtocol,
+)
+
+_LOGGER = logging.getLogger(__name__)
+
+A01_VERSION = b"A01"
+
+
+def _no_encode(value: Any) -> Any:
+    return value
+
+
+def encode_mqtt_payload(
+    data: dict[RoborockDyadDataProtocol, Any]
+    | dict[RoborockZeoProtocol, Any]
+    | dict[RoborockDyadDataProtocol | RoborockZeoProtocol, Any],
+    value_encoder: Callable[[Any], Any] | None = None,
+) -> RoborockMessage:
+    """Encode payload for A01 commands over MQTT.
+
+    Args:
+        data: The data to encode.
+        value_encoder: A function to encode the values of the dictionary.
+
+    Returns:
+        RoborockMessage: The encoded message.
+    """
+    if value_encoder is None:
+        value_encoder = _no_encode
+    dps_data = {"dps": {key: value_encoder(value) for key, value in data.items()}}
+    payload = pad(json.dumps(dps_data).encode("utf-8"), AES.block_size)
+    return RoborockMessage(
+        protocol=RoborockMessageProtocol.RPC_REQUEST,
+        version=A01_VERSION,
+        payload=payload,
+    )
+
+
+def decode_rpc_response(message: RoborockMessage) -> dict[int, Any]:
+    """Decode a V1 RPC_RESPONSE message."""
+    if not message.payload:
+        raise RoborockException("Invalid A01 message format: missing payload")
+    try:
+        unpadded = unpad(message.payload, AES.block_size)
+    except ValueError as err:
+        raise RoborockException(f"Unable to unpad A01 payload: {err}")
+
+    try:
+        payload = json.loads(unpadded.decode())
+    except (json.JSONDecodeError, TypeError) as e:
+        raise RoborockException(f"Invalid A01 message payload: {e} for {message.payload!r}") from e
+
+    datapoints = payload.get("dps", {})
+    if not isinstance(datapoints, dict):
+        raise RoborockException(f"Invalid A01 message format: 'dps' should be a dictionary for {message.payload!r}")
+    try:
+        return {int(key): value for key, value in datapoints.items()}
+    except ValueError:
+        raise RoborockException(f"Invalid A01 message format: 'dps' key should be an integer for {message.payload!r}")

roborock_cli/_vendor/roborock/protocols/b01_q10_protocol.py

@@ -0,0 +1,87 @@
+"""Roborock B01 Protocol encoding and decoding."""
+
+import json
+import logging
+from typing import Any
+
+from roborock_cli._vendor.roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP
+from roborock_cli._vendor.roborock.exceptions import RoborockException
+from roborock_cli._vendor.roborock.roborock_message import (
+    RoborockMessage,
+    RoborockMessageProtocol,
+)
+
+_LOGGER = logging.getLogger(__name__)
+
+B01_VERSION = b"B01"
+ParamsType = list | dict | int | None
+
+
+def encode_mqtt_payload(command: B01_Q10_DP, params: ParamsType) -> RoborockMessage:
+    """Encode payload for B01 Q10 commands over MQTT.
+
+    This does not perform any special encoding for the command parameters and expects
+    them to already be in a request specific format.
+    """
+    dps_data = {
+        "dps": {
+            # Important: some commands use falsy values so only default to `{}` when params is actually None.
+            command.code: params if params is not None else {},
+        }
+    }
+    return RoborockMessage(
+        protocol=RoborockMessageProtocol.RPC_REQUEST,
+        version=B01_VERSION,
+        payload=json.dumps(dps_data).encode("utf-8"),
+    )
+
+
+def _convert_datapoints(datapoints: dict[str, Any], message: RoborockMessage) -> dict[B01_Q10_DP, Any]:
+    """Convert the 'dps' dictionary keys from strings to B01_Q10_DP enums."""
+    result: dict[B01_Q10_DP, Any] = {}
+    for key, value in datapoints.items():
+        try:
+            code = int(key)
+        except ValueError as e:
+            raise ValueError(f"dps key is not a valid integer: {e} for {message.payload!r}") from e
+        if (dps := B01_Q10_DP.from_code_optional(code)) is not None:
+            result[dps] = value
+    return result
+
+
+def decode_rpc_response(message: RoborockMessage) -> dict[B01_Q10_DP, Any]:
+    """Decode a B01 Q10 RPC_RESPONSE message.
+
+    This does not perform any special decoding for the response body, but does
+    convert the 'dps' keys from strings to B01_Q10_DP enums.
+    """
+    if not message.payload:
+        raise RoborockException("Invalid B01 message format: missing payload")
+    try:
+        payload = json.loads(message.payload.decode())
+    except (json.JSONDecodeError, UnicodeDecodeError) as e:
+        raise RoborockException(f"Invalid B01 json payload: {e} for {message.payload!r}") from e
+
+    if (datapoints := payload.get("dps")) is None:
+        raise RoborockException(f"Invalid B01 json payload: missing 'dps' for {message.payload!r}")
+    if not isinstance(datapoints, dict):
+        raise RoborockException(f"Invalid B01 message format: 'dps' should be a dictionary for {message.payload!r}")
+
+    try:
+        result = _convert_datapoints(datapoints, message)
+    except ValueError as e:
+        raise RoborockException(f"Invalid B01 message format: {e}") from e
+
+    # The COMMON response contains nested datapoints need conversion. To simplify
+    # response handling at higher levels we flatten these into the main result.
+    if B01_Q10_DP.COMMON in result:
+        common_result = result.pop(B01_Q10_DP.COMMON)
+        if not isinstance(common_result, dict):
+            raise RoborockException(f"Invalid dpCommon format: expected dict, got {type(common_result).__name__}")
+        try:
+            common_dps_result = _convert_datapoints(common_result, message)
+        except ValueError as e:
+            raise RoborockException(f"Invalid dpCommon format: {e}") from e
+        result.update(common_dps_result)
+
+    return result

roborock_cli/_vendor/roborock/protocols/b01_q7_protocol.py

@@ -0,0 +1,81 @@
+"""Roborock B01 Protocol encoding and decoding."""
+
+import json
+import logging
+from dataclasses import dataclass, field
+from typing import Any
+
+from Crypto.Cipher import AES
+from Crypto.Util.Padding import pad, unpad
+
+from roborock_cli._vendor.roborock import RoborockB01Q7Methods
+from roborock_cli._vendor.roborock.exceptions import RoborockException
+from roborock_cli._vendor.roborock.roborock_message import (
+    RoborockMessage,
+    RoborockMessageProtocol,
+)
+from roborock_cli._vendor.roborock.util import get_next_int
+
+_LOGGER = logging.getLogger(__name__)
+
+B01_VERSION = b"B01"
+CommandType = RoborockB01Q7Methods | str
+ParamsType = list | dict | int | None
+
+
+@dataclass
+class Q7RequestMessage:
+    """Data class for B01 Q7 request message."""
+
+    dps: int
+    command: CommandType
+    params: ParamsType
+    msg_id: int = field(default_factory=lambda: get_next_int(100000000000, 999999999999))
+
+    def to_dps_value(self) -> dict[int, Any]:
+        """Return the 'dps' payload dictionary."""
+        return {
+            self.dps: {
+                "method": str(self.command),
+                "msgId": str(self.msg_id),
+                # Important: some B01 methods use an empty object `{}` (not `[]`) for
+                # "no params", and some setters legitimately send `0` which is falsy.
+                # Only default to `[]` when params is actually None.
+                "params": self.params if self.params is not None else [],
+            }
+        }
+
+
+def encode_mqtt_payload(request: Q7RequestMessage) -> RoborockMessage:
+    """Encode payload for B01 commands over MQTT."""
+    dps_data = {"dps": request.to_dps_value()}
+    payload = pad(json.dumps(dps_data).encode("utf-8"), AES.block_size)
+    return RoborockMessage(
+        protocol=RoborockMessageProtocol.RPC_REQUEST,
+        version=B01_VERSION,
+        payload=payload,
+    )
+
+
+def decode_rpc_response(message: RoborockMessage) -> dict[int, Any]:
+    """Decode a B01 RPC_RESPONSE message."""
+    if not message.payload:
+        raise RoborockException("Invalid B01 message format: missing payload")
+    try:
+        unpadded = unpad(message.payload, AES.block_size)
+    except ValueError:
+        # It would be better to fail down the line.
+        unpadded = message.payload
+
+    try:
+        payload = json.loads(unpadded.decode())
+    except (json.JSONDecodeError, TypeError) as e:
+        raise RoborockException(f"Invalid B01 message payload: {e} for {message.payload!r}") from e
+
+    datapoints = payload.get("dps", {})
+    if not isinstance(datapoints, dict):
+        raise RoborockException(f"Invalid B01 message format: 'dps' should be a dictionary for {message.payload!r}")
+    try:
+        return {int(key): value for key, value in datapoints.items()}
+    except ValueError:
+        raise RoborockException(f"Invalid B01 message format: 'dps' key should be an integer for {message.payload!r}")

roborock_cli/_vendor/roborock/protocols/v1_protocol.py

@@ -0,0 +1,271 @@
+"""Roborock V1 Protocol Encoder."""
+
+from __future__ import annotations
+
+import base64
+import json
+import logging
+import secrets
+import struct
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from enum import StrEnum
+from typing import Any, Protocol, TypeVar, overload
+
+from roborock_cli._vendor.roborock.data import RoborockBase, RRiot
+from roborock_cli._vendor.roborock.exceptions import RoborockException, RoborockInvalidStatus, RoborockUnsupportedFeature
+from roborock_cli._vendor.roborock.protocol import Utils
+from roborock_cli._vendor.roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
+from roborock_cli._vendor.roborock.roborock_typing import RoborockCommand
+from roborock_cli._vendor.roborock.util import get_next_int, get_timestamp
+
+_LOGGER = logging.getLogger(__name__)
+
+__all__ = [
+    "SecurityData",
+    "create_security_data",
+    "decode_rpc_response",
+    "V1RpcChannel",
+]
+
+CommandType = RoborockCommand | str
+ParamsType = list | dict | int | None
+
+
+class LocalProtocolVersion(StrEnum):
+    """Supported local protocol versions. Different from vacuum protocol versions."""
+
+    L01 = "L01"
+    V1 = "1.0"
+
+
+@dataclass(frozen=True, kw_only=True)
+class SecurityData:
+    """Security data included in the request for some V1 commands."""
+
+    endpoint: str
+    nonce: bytes
+
+    def to_dict(self) -> dict[str, Any]:
+        """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."""
+    nonce = secrets.token_bytes(16)
+    endpoint = base64.b64encode(Utils.md5(rriot.k.encode())[8:14]).decode()
+    return SecurityData(endpoint=endpoint, nonce=nonce)
+
+
+@dataclass
+class RequestMessage:
+    """Data structure for v1 RoborockMessage payloads."""
+
+    method: RoborockCommand | str
+    params: ParamsType
+    timestamp: int = field(default_factory=lambda: get_timestamp())
+    request_id: int = field(default_factory=lambda: get_next_int(10000, 32767))
+
+    def encode_message(
+        self,
+        protocol: RoborockMessageProtocol,
+        security_data: SecurityData | None = None,
+        version: LocalProtocolVersion = LocalProtocolVersion.V1,
+    ) -> RoborockMessage:
+        """Convert the request message to a RoborockMessage."""
+        return RoborockMessage(
+            timestamp=self.timestamp,
+            protocol=protocol,
+            payload=self._as_payload(security_data=security_data),
+            version=version.value.encode(),
+        )
+
+    def _as_payload(self, security_data: SecurityData | None) -> bytes:
+        """Convert the request arguments to a dictionary."""
+        inner = {
+            "id": self.request_id,
+            "method": self.method,
+            "params": self.params or [],
+            **(security_data.to_dict() if security_data else {}),
+        }
+        return bytes(
+            json.dumps(
+                {
+                    "dps": {"101": json.dumps(inner, separators=(",", ":"))},
+                    "t": self.timestamp,
+                },
+                separators=(",", ":"),
+            ).encode()
+        )
+
+
+ResponseData = dict[str, Any] | list | int
+
+# V1 RPC error code mappings to specific exception types
+_V1_ERROR_CODE_EXCEPTIONS: dict[int, type[RoborockException]] = {
+    -10007: RoborockInvalidStatus,  # "invalid status" - device action locked
+}
+
+
+def _create_api_error(error: Any) -> RoborockException:
+    """Create an appropriate exception for a V1 RPC error response.
+
+    Maps known error codes to specific exception types for easier handling
+    at higher levels.
+    """
+    if isinstance(error, dict):
+        code = error.get("code")
+        if isinstance(code, int) and (exc_type := _V1_ERROR_CODE_EXCEPTIONS.get(code)):
+            return exc_type(error)
+    return RoborockException(error)
+
+
+@dataclass(kw_only=True, frozen=True)
+class ResponseMessage:
+    """Data structure for v1 RoborockMessage responses."""
+
+    request_id: int | None
+    """The request ID of the response."""
+
+    data: ResponseData
+    """The data of the response, where the type depends on the command."""
+
+    api_error: RoborockException | None = None
+    """The API error message of the response if any."""
+
+
+def decode_rpc_response(message: RoborockMessage) -> ResponseMessage:
+    """Decode a V1 RPC_RESPONSE message.
+
+    This will raise a RoborockException if the message cannot be parsed. A
+    response object will be returned even if there is an error in the
+    response, as long as we can extract the request ID. This is so we can
+    associate an API response with a request even if there was an error.
+    """
+    if not message.payload:
+        return ResponseMessage(request_id=message.seq, data={})
+    try:
+        payload = json.loads(message.payload.decode())
+    except (json.JSONDecodeError, TypeError, UnicodeDecodeError) as e:
+        raise RoborockException(f"Invalid V1 message payload: {e} for {message.payload!r}") from e
+
+    _LOGGER.debug("Decoded V1 message payload: %s", payload)
+    datapoints = payload.get("dps", {})
+    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(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 '{RoborockMessageProtocol.RPC_RESPONSE}': {e} for {message.payload!r}"
+        ) from e
+
+    request_id: int | None = data_point_response.get("id")
+    api_error: RoborockException | None = None
+    if error := data_point_response.get("error"):
+        api_error = _create_api_error(error)
+
+    if (result := data_point_response.get("result")) is None:
+        # Some firmware versions return an error-only response (no "result" key).
+        # Preserve that error instead of overwriting it with a parsing exception.
+        if api_error is None:
+            api_error = RoborockException(
+                f"Invalid V1 message format: missing 'result' in data point for {message.payload!r}"
+            )
+        result = {}
+    else:
+        _LOGGER.debug("Decoded V1 message result: %s", result)
+        if isinstance(result, str):
+            if result == "unknown_method":
+                api_error = RoborockUnsupportedFeature("The method called is not recognized by the device.")
+            elif result != "ok":
+                api_error = RoborockException(f"Unexpected API Result: {result}")
+            result = {}
+        if not isinstance(result, dict | list | int):
+            # If we already have an API error, prefer returning a response object
+            # rather than failing to decode the message entirely.
+            if api_error is None:
+                raise RoborockException(
+                    f"Invalid V1 message format: 'result' was unexpected type {type(result)}. {message.payload!r}"
+                )
+            result = {}
+
+    if not request_id and api_error:
+        raise api_error
+    return ResponseMessage(request_id=request_id, data=result, api_error=api_error)
+
+
+@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 | None]:
+    """Create a decoder for V1 map response messages."""
+
+    def _decode_map_response(message: RoborockMessage) -> MapResponse | None:
+        """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):
+            _LOGGER.debug("Received map response not requested by this device, ignoring.")
+            return None
+        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
+
+
+_T = TypeVar("_T", bound=RoborockBase)
+
+
+class V1RpcChannel(Protocol):
+    """Protocol for V1 RPC channels.
+
+    This is a wrapper around a raw channel that provides a high-level interface
+    for sending commands and receiving responses.
+    """
+
+    @overload
+    async def send_command(
+        self,
+        method: CommandType,
+        *,
+        params: ParamsType = None,
+    ) -> Any:
+        """Send a command and return a decoded response."""
+        ...
+
+    @overload
+    async def send_command(
+        self,
+        method: CommandType,
+        *,
+        response_type: type[_T],
+        params: ParamsType = None,
+    ) -> _T:
+        """Send a command and return a parsed response RoborockBase type."""
+        ...