roborock-cli 0.1.1__py3-none-any.whl

roborock_cli/_vendor/roborock/devices/device_manager.py

@@ -0,0 +1,269 @@
+"""Module for discovering Roborock devices."""
+
+import asyncio
+import enum
+import logging
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import Any
+
+import aiohttp
+
+from roborock_cli._vendor.roborock.data import (
+    HomeData,
+    HomeDataDevice,
+    HomeDataProduct,
+    UserData,
+)
+from roborock_cli._vendor.roborock.devices.device import DeviceReadyCallback, RoborockDevice
+from roborock_cli._vendor.roborock.diagnostics import Diagnostics, redact_device_data
+from roborock_cli._vendor.roborock.exceptions import RoborockException
+from roborock_cli._vendor.roborock.map.map_parser import MapParserConfig
+from roborock_cli._vendor.roborock.mqtt.roborock_session import create_lazy_mqtt_session
+from roborock_cli._vendor.roborock.mqtt.session import MqttSession, SessionUnauthorizedHook
+from roborock_cli._vendor.roborock.protocol import create_mqtt_params
+from roborock_cli._vendor.roborock.web_api import RoborockApiClient, UserWebApiClient
+
+from .cache import Cache, DeviceCache, NoCache
+from .rpc.v1_channel import create_v1_channel
+from .traits import Trait, a01, b01, v1
+from .transport.channel import Channel
+from .transport.mqtt_channel import create_mqtt_channel
+
+_LOGGER = logging.getLogger(__name__)
+
+__all__ = [
+    "create_device_manager",
+    "UserParams",
+    "DeviceManager",
+]
+
+
+DeviceCreator = Callable[[HomeData, HomeDataDevice, HomeDataProduct], RoborockDevice]
+
+
+class DeviceVersion(enum.StrEnum):
+    """Enum for device versions."""
+
+    V1 = "1.0"
+    A01 = "A01"
+    B01 = "B01"
+    UNKNOWN = "unknown"
+
+
+class UnsupportedDeviceError(RoborockException):
+    """Exception raised when a device is unsupported."""
+
+
+class DeviceManager:
+    """Central manager for Roborock device discovery and connections."""
+
+    def __init__(
+        self,
+        web_api: UserWebApiClient,
+        device_creator: DeviceCreator,
+        mqtt_session: MqttSession,
+        cache: Cache,
+        diagnostics: Diagnostics,
+    ) -> None:
+        """Initialize the DeviceManager with user data and optional cache storage.
+
+        This takes ownership of the MQTT session and will close it when the manager is closed.
+        """
+        self._web_api = web_api
+        self._cache = cache
+        self._device_creator = device_creator
+        self._devices: dict[str, RoborockDevice] = {}
+        self._mqtt_session = mqtt_session
+        self._diagnostics = diagnostics
+        self._home_data: HomeData | None = None
+
+    async def discover_devices(self, prefer_cache: bool = True) -> list[RoborockDevice]:
+        """Discover all devices for the logged-in user."""
+        self._diagnostics.increment("discover_devices")
+        cache_data = await self._cache.get()
+        if not cache_data.home_data or not prefer_cache:
+            _LOGGER.debug("Fetching home data (prefer_cache=%s)", prefer_cache)
+            self._diagnostics.increment("fetch_home_data")
+            try:
+                cache_data.home_data = await self._web_api.get_home_data()
+            except RoborockException as ex:
+                if not cache_data.home_data:
+                    raise
+                _LOGGER.debug("Failed to fetch home data, using cached data: %s", ex)
+            await self._cache.set(cache_data)
+        self._home_data = cache_data.home_data
+
+        device_products = self._home_data.device_products
+        _LOGGER.debug("Discovered %d devices", len(device_products))
+
+        # These are connected serially to avoid overwhelming the MQTT broker
+        new_devices = {}
+        start_tasks = []
+        supported_devices_counter = self._diagnostics.subkey("supported_devices")
+        unsupported_devices_counter = self._diagnostics.subkey("unsupported_devices")
+        for duid, (device, product) in device_products.items():
+            _LOGGER.debug("[%s] Discovered device %s %s", duid, product.summary_info(), device.summary_info())
+            if duid in self._devices:
+                continue
+            try:
+                new_device = self._device_creator(self._home_data, device, product)
+            except UnsupportedDeviceError:
+                _LOGGER.info("Skipping unsupported device %s %s", product.summary_info(), device.summary_info())
+                unsupported_devices_counter.increment(device.pv or "unknown")
+                continue
+            supported_devices_counter.increment(device.pv or "unknown")
+            start_tasks.append(new_device.start_connect())
+            new_devices[duid] = new_device
+
+        self._devices.update(new_devices)
+        await asyncio.gather(*start_tasks)
+        return list(self._devices.values())
+
+    async def get_device(self, duid: str) -> RoborockDevice | None:
+        """Get a specific device by DUID."""
+        return self._devices.get(duid)
+
+    async def get_devices(self) -> list[RoborockDevice]:
+        """Get all discovered devices."""
+        return list(self._devices.values())
+
+    async def close(self) -> None:
+        """Close all MQTT connections and clean up resources."""
+        tasks = [device.close() for device in self._devices.values()]
+        self._devices.clear()
+        tasks.append(self._mqtt_session.close())
+        await asyncio.gather(*tasks)
+
+    def diagnostic_data(self) -> Mapping[str, Any]:
+        """Return diagnostics information about the device manager."""
+        return {
+            "home_data": redact_device_data(self._home_data.as_dict()) if self._home_data else None,
+            "devices": [device.diagnostic_data() for device in self._devices.values()],
+            "diagnostics": self._diagnostics.as_dict(),
+        }
+
+
+@dataclass
+class UserParams:
+    """Parameters for creating a new session with Roborock devices.
+
+    These parameters include the username, user data for authentication,
+    and an optional base URL for the Roborock API. The `user_data` and `base_url`
+    parameters are obtained from `RoborockApiClient` during the login process.
+    """
+
+    username: str
+    """The username (email) used for logging in."""
+
+    user_data: UserData
+    """This is the user data containing authentication information."""
+
+    base_url: str | None = None
+    """Optional base URL for the Roborock API.
+
+    This is used to speed up connection times by avoiding the need to
+    discover the API base URL each time. If not provided, the API client
+    will attempt to discover it automatically which may take multiple requests.
+    """
+
+
+def create_web_api_wrapper(
+    user_params: UserParams,
+    *,
+    cache: Cache | None = None,
+    session: aiohttp.ClientSession | None = None,
+) -> UserWebApiClient:
+    """Create a home data API wrapper from an existing API client."""
+
+    # Note: This will auto discover the API base URL. This can be improved
+    # by caching this next to `UserData` if needed to avoid unnecessary API calls.
+    client = RoborockApiClient(username=user_params.username, base_url=user_params.base_url, session=session)
+
+    return UserWebApiClient(client, user_params.user_data)
+
+
+async def create_device_manager(
+    user_params: UserParams,
+    *,
+    cache: Cache | None = None,
+    map_parser_config: MapParserConfig | None = None,
+    session: aiohttp.ClientSession | None = None,
+    ready_callback: DeviceReadyCallback | None = None,
+    mqtt_session_unauthorized_hook: SessionUnauthorizedHook | None = None,
+    prefer_cache: bool = True,
+) -> DeviceManager:
+    """Convenience function to create and initialize a DeviceManager.
+
+    Args:
+        user_params: Parameters for creating the user session.
+        cache: Optional cache implementation to use for caching device data.
+        map_parser_config: Optional configuration for parsing maps.
+        session: Optional aiohttp ClientSession to use for HTTP requests.
+        ready_callback: Optional callback to be notified when a device is ready.
+        mqtt_session_unauthorized_hook: Optional hook for MQTT session unauthorized
+          events which may indicate rate limiting or revoked credentials. The
+          caller may use this to refresh authentication tokens as needed.
+        prefer_cache: Whether to prefer cached device data over always fetching it from the API.
+
+    Returns:
+        An initialized DeviceManager with discovered devices.
+    """
+    if cache is None:
+        cache = NoCache()
+
+    web_api = create_web_api_wrapper(user_params, session=session, cache=cache)
+    user_data = user_params.user_data
+
+    diagnostics = Diagnostics()
+
+    mqtt_params = create_mqtt_params(user_data.rriot)
+    mqtt_params.diagnostics = diagnostics.subkey("mqtt_session")
+    mqtt_params.unauthorized_hook = mqtt_session_unauthorized_hook
+    mqtt_session = await create_lazy_mqtt_session(mqtt_params)
+
+    def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDataProduct) -> RoborockDevice:
+        channel: Channel
+        trait: Trait
+        device_cache: DeviceCache = DeviceCache(device.duid, cache)
+        match device.pv:
+            case DeviceVersion.V1:
+                channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device, device_cache)
+                trait = v1.create(
+                    device.duid,
+                    product,
+                    home_data,
+                    channel.rpc_channel,
+                    channel.mqtt_rpc_channel,
+                    channel.map_rpc_channel,
+                    web_api,
+                    device_cache=device_cache,
+                    map_parser_config=map_parser_config,
+                    region=user_data.region,
+                )
+            case DeviceVersion.A01:
+                channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
+                trait = a01.create(product, channel)
+            case DeviceVersion.B01:
+                channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
+                model_part = product.model.split(".")[-1]
+                if "ss" in model_part:
+                    trait = b01.q10.create(channel)
+                elif "sc" in model_part:
+                    # Q7 devices start with 'sc' in their model naming.
+                    trait = b01.q7.create(channel)
+                else:
+                    raise UnsupportedDeviceError(f"Device {device.name} has unsupported B01 model: {product.model}")
+            case _:
+                raise UnsupportedDeviceError(
+                    f"Device {device.name} has unsupported version {device.pv} {product.model}"
+                )
+
+        dev = RoborockDevice(device, product, channel, trait)
+        if ready_callback:
+            dev.add_ready_callback(ready_callback)
+        return dev
+
+    manager = DeviceManager(web_api, device_creator, mqtt_session=mqtt_session, cache=cache, diagnostics=diagnostics)
+    await manager.discover_devices(prefer_cache)
+    return manager

roborock_cli/_vendor/roborock/devices/file_cache.py

@@ -0,0 +1,79 @@
+"""This module implements a file-backed cache for device information.
+
+This module provides a `FileCache` class that implements the `Cache` protocol
+to store and retrieve cached device information from a file on disk. This allows
+persistent caching of device data across application restarts.
+"""
+
+import asyncio
+import pathlib
+import pickle
+from collections.abc import Callable
+from typing import Any
+
+from .cache import Cache, CacheData
+
+
+class FileCache(Cache):
+    """File backed cache implementation."""
+
+    def __init__(
+        self,
+        file_path: pathlib.Path,
+        init_fn: Callable[[], CacheData] = CacheData,
+        serialize_fn: Callable[[Any], bytes] = pickle.dumps,
+        deserialize_fn: Callable[[bytes], Any] = pickle.loads,
+    ) -> None:
+        """Initialize the file cache with the given file path."""
+        self._init_fn = init_fn
+        self._file_path = file_path
+        self._cache_data: CacheData | None = None
+        self._serialize_fn = serialize_fn
+        self._deserialize_fn = deserialize_fn
+
+    async def get(self) -> CacheData:
+        """Get cached value."""
+        if self._cache_data is not None:
+            return self._cache_data
+        data = await load_value(self._file_path, self._deserialize_fn)
+        if data is not None and not isinstance(data, CacheData):
+            raise TypeError(f"Invalid cache data loaded from {self._file_path}")
+
+        self._cache_data = data or self._init_fn()
+        return self._cache_data
+
+    async def set(self, value: CacheData) -> None:  # type: ignore[override]
+        """Set value in the cache."""
+        self._cache_data = value
+
+    async def flush(self) -> None:
+        """Flush the cache to disk."""
+        if self._cache_data is None:
+            return
+        await store_value(self._file_path, self._cache_data, self._serialize_fn)
+
+
+async def store_value(file_path: pathlib.Path, value: Any, serialize_fn: Callable[[Any], bytes] = pickle.dumps) -> None:
+    """Store a value to the given file path."""
+
+    def _store_to_disk(file_path: pathlib.Path, value: Any) -> None:
+        with open(file_path, "wb") as f:
+            data = serialize_fn(value)
+            f.write(data)
+
+    await asyncio.to_thread(_store_to_disk, file_path, value)
+
+
+async def load_value(file_path: pathlib.Path, deserialize_fn: Callable[[bytes], Any] = pickle.loads) -> Any | None:
+    """Load a value from the given file path."""
+
+    def _load_from_disk(file_path: pathlib.Path) -> Any | None:
+        if not file_path.exists():
+            return None
+        with open(file_path, "rb") as f:
+            data = f.read()
+            if not data:  # 空文件视为不存在
+                return None
+            return deserialize_fn(data)
+
+    return await asyncio.to_thread(_load_from_disk, file_path)

roborock_cli/_vendor/roborock/devices/rpc/__init__.py

@@ -0,0 +1,14 @@
+"""Module for sending device specific commands to Roborock devices.
+
+This module provides a application-level interface for sending commands to Roborock
+devices. These modules can be used by traits (higher level APIs) to send commands.
+
+Each module may contain details that are common across all traits, and may depend
+on the transport level modules (e.g. MQTT, Local device) for issuing the
+commands.
+
+The lowest level protocol encoding is handled in `roborock.protocols` which
+have no dependencies on the transport level modules.
+"""
+
+__all__: list[str] = []

roborock_cli/_vendor/roborock/devices/rpc/a01_channel.py

@@ -0,0 +1,94 @@
+"""Thin wrapper around the MQTT channel for Roborock A01 devices."""
+
+import asyncio
+import logging
+from collections.abc import Callable
+from typing import Any, overload
+
+from roborock_cli._vendor.roborock.devices.transport.mqtt_channel import MqttChannel
+from roborock_cli._vendor.roborock.exceptions import RoborockException
+from roborock_cli._vendor.roborock.protocols.a01_protocol import (
+    decode_rpc_response,
+    encode_mqtt_payload,
+)
+from roborock_cli._vendor.roborock.roborock_message import (
+    RoborockDyadDataProtocol,
+    RoborockMessage,
+    RoborockZeoProtocol,
+)
+
+_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],
+    value_encoder: Callable[[Any], Any] | None = None,
+) -> dict[RoborockDyadDataProtocol, Any]: ...
+
+
+@overload
+async def send_decoded_command(
+    mqtt_channel: MqttChannel,
+    params: dict[RoborockZeoProtocol, Any],
+    value_encoder: Callable[[Any], Any] | None = None,
+) -> dict[RoborockZeoProtocol, Any]: ...
+
+
+async def send_decoded_command(
+    mqtt_channel: MqttChannel,
+    params: dict[RoborockDyadDataProtocol, Any] | dict[RoborockZeoProtocol, Any],
+    value_encoder: Callable[[Any], Any] | None = None,
+) -> 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, value_encoder)
+
+    # 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]

roborock_cli/_vendor/roborock/devices/rpc/b01_q10_channel.py

@@ -0,0 +1,57 @@
+"""Thin wrapper around the MQTT channel for Roborock B01 Q10 devices."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import AsyncGenerator
+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.devices.transport.mqtt_channel import MqttChannel
+from roborock_cli._vendor.roborock.exceptions import RoborockException
+from roborock_cli._vendor.roborock.protocols.b01_q10_protocol import (
+    ParamsType,
+    decode_rpc_response,
+    encode_mqtt_payload,
+)
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def stream_decoded_responses(
+    mqtt_channel: MqttChannel,
+) -> AsyncGenerator[dict[B01_Q10_DP, Any], None]:
+    """Stream decoded DPS messages received via MQTT."""
+
+    async for response_message in mqtt_channel.subscribe_stream():
+        try:
+            decoded_dps = decode_rpc_response(response_message)
+        except RoborockException as ex:
+            _LOGGER.debug(
+                "Failed to decode B01 Q10 RPC response: %s: %s",
+                response_message,
+                ex,
+            )
+            continue
+        yield decoded_dps
+
+
+async def send_command(
+    mqtt_channel: MqttChannel,
+    command: B01_Q10_DP,
+    params: ParamsType,
+) -> None:
+    """Send a command on the MQTT channel, without waiting for a response"""
+    _LOGGER.debug("Sending B01 MQTT command: cmd=%s params=%s", command, params)
+    roborock_message = encode_mqtt_payload(command, params)
+    _LOGGER.debug("Sending MQTT message: %s", roborock_message)
+    try:
+        await mqtt_channel.publish(roborock_message)
+    except RoborockException as ex:
+        _LOGGER.debug(
+            "Error sending B01 decoded command (method=%s params=%s): %s",
+            command,
+            params,
+            ex,
+        )
+        raise

roborock_cli/_vendor/roborock/devices/rpc/b01_q7_channel.py

@@ -0,0 +1,101 @@
+"""Thin wrapper around the MQTT channel for Roborock B01 Q7 devices."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from typing import Any
+
+from roborock_cli._vendor.roborock.devices.transport.mqtt_channel import MqttChannel
+from roborock_cli._vendor.roborock.exceptions import RoborockException
+from roborock_cli._vendor.roborock.protocols.b01_q7_protocol import (
+    Q7RequestMessage,
+    decode_rpc_response,
+    encode_mqtt_payload,
+)
+from roborock_cli._vendor.roborock.roborock_message import RoborockMessage
+
+_LOGGER = logging.getLogger(__name__)
+_TIMEOUT = 10.0
+
+
+async def send_decoded_command(
+    mqtt_channel: MqttChannel,
+    request_message: Q7RequestMessage,
+) -> dict[str, Any] | None:
+    """Send a command on the MQTT channel and get a decoded response."""
+    _LOGGER.debug("Sending B01 MQTT command: %s", request_message)
+    roborock_message = encode_mqtt_payload(request_message)
+    future: asyncio.Future[Any] = asyncio.get_running_loop().create_future()
+
+    def find_response(response_message: RoborockMessage) -> None:
+        """Handle incoming messages and resolve the future."""
+        try:
+            decoded_dps = decode_rpc_response(response_message)
+        except RoborockException as ex:
+            _LOGGER.debug(
+                "Failed to decode B01 RPC response (expecting method=%s msg_id=%s): %s: %s",
+                request_message.command,
+                request_message.msg_id,
+                response_message,
+                ex,
+            )
+            return
+        for dps_value in decoded_dps.values():
+            # valid responses are JSON strings wrapped in the dps value
+            if not isinstance(dps_value, str):
+                _LOGGER.debug("Received unexpected response: %s", dps_value)
+                continue
+
+            try:
+                inner = json.loads(dps_value)
+            except (json.JSONDecodeError, TypeError):
+                _LOGGER.debug("Received unexpected response: %s", dps_value)
+                continue
+            if isinstance(inner, dict) and inner.get("msgId") == str(request_message.msg_id):
+                _LOGGER.debug("Received query response: %s", inner)
+                # Check for error code (0 = success, non-zero = error)
+                code = inner.get("code", 0)
+                if code != 0:
+                    error_msg = f"B01 command failed with code {code} ({request_message})"
+                    _LOGGER.debug("B01 error response: %s", error_msg)
+                    if not future.done():
+                        future.set_exception(RoborockException(error_msg))
+                    return
+                data = inner.get("data")
+                # All get commands should be dicts
+                if request_message.command.endswith(".get") and not isinstance(data, dict):
+                    if not future.done():
+                        future.set_exception(
+                            RoborockException(f"Unexpected data type for response {data} ({request_message})")
+                        )
+                    return
+                if not future.done():
+                    future.set_result(data)
+
+    unsub = await mqtt_channel.subscribe(find_response)
+
+    _LOGGER.debug("Sending MQTT message: %s", roborock_message)
+    try:
+        await mqtt_channel.publish(roborock_message)
+        return await asyncio.wait_for(future, timeout=_TIMEOUT)
+    except TimeoutError as ex:
+        raise RoborockException(f"B01 command timed out after {_TIMEOUT}s ({request_message})") from ex
+    except RoborockException as ex:
+        _LOGGER.warning(
+            "Error sending B01 decoded command (%ss): %s",
+            request_message,
+            ex,
+        )
+        raise
+
+    except Exception as ex:
+        _LOGGER.exception(
+            "Error sending B01 decoded command (%ss): %s",
+            request_message,
+            ex,
+        )
+        raise
+    finally:
+        unsub()