roborock-cli 0.1.1__py3-none-any.whl

roborock_cli/_vendor/roborock/devices/transport/local_channel.py

@@ -0,0 +1,295 @@
+"""Module for communicating with Roborock devices over a local network."""
+
+import asyncio
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from roborock_cli._vendor.roborock.callbacks import CallbackList, decoder_callback
+from roborock_cli._vendor.roborock.exceptions import RoborockConnectionException, RoborockException
+from roborock_cli._vendor.roborock.protocol import create_local_decoder, create_local_encoder
+from roborock_cli._vendor.roborock.protocols.v1_protocol import LocalProtocolVersion
+from roborock_cli._vendor.roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
+from roborock_cli._vendor.roborock.util import RoborockLoggerAdapter, get_next_int
+
+from .channel import Channel
+
+_LOGGER = logging.getLogger(__name__)
+_PORT = 58867
+_TIMEOUT = 5.0
+_PING_INTERVAL = 10
+
+
+@dataclass
+class LocalChannelParams:
+    """Parameters for local channel encoder/decoder."""
+
+    local_key: str
+    connect_nonce: int
+    ack_nonce: int | None
+
+
+@dataclass
+class _LocalProtocol(asyncio.Protocol):
+    """Callbacks for the Roborock local client transport."""
+
+    messages_cb: Callable[[bytes], None]
+    connection_lost_cb: Callable[[Exception | None], None]
+
+    def data_received(self, data: bytes) -> None:
+        """Called when data is received from the transport."""
+        self.messages_cb(data)
+
+    def connection_lost(self, exc: Exception | None) -> None:
+        """Called when the transport connection is lost."""
+        self.connection_lost_cb(exc)
+
+
+def get_running_loop() -> asyncio.AbstractEventLoop:
+    """Get the running event loop, extracted for mocking purposes."""
+    return asyncio.get_running_loop()
+
+
+class LocalChannel(Channel):
+    """Simple RPC-style channel for communicating with a device over a local network.
+
+    Handles request/response correlation and timeouts, but leaves message
+    format most parsing to higher-level components.
+    """
+
+    def __init__(self, host: str, local_key: str, device_uid: str) -> None:
+        self._host = host
+        self._logger = RoborockLoggerAdapter(duid=device_uid, logger=_LOGGER)
+        self._transport: asyncio.Transport | None = None
+        self._protocol: _LocalProtocol | None = None
+        self._subscribers: CallbackList[RoborockMessage] = CallbackList(self._logger)
+        self._is_connected = False
+        self._local_protocol_version: LocalProtocolVersion | None = None
+        self._keep_alive_task: asyncio.Task[None] | None = None
+        self._update_encoder_decoder(
+            LocalChannelParams(local_key=local_key, connect_nonce=get_next_int(10000, 32767), ack_nonce=None)
+        )
+
+    def _update_encoder_decoder(self, params: LocalChannelParams) -> None:
+        """Update the encoder and decoder with new parameters.
+
+        This is invoked once with an initial set of values used for protocol
+        negotiation. Once negotiation completes, it is updated again to set the
+        correct nonces for the follow up communications and updates the encoder
+        and decoder functions accordingly.
+        """
+        self._params = params
+        self._encoder = create_local_encoder(
+            local_key=params.local_key, connect_nonce=params.connect_nonce, ack_nonce=params.ack_nonce
+        )
+        self._decoder = create_local_decoder(
+            local_key=params.local_key, connect_nonce=params.connect_nonce, ack_nonce=params.ack_nonce
+        )
+        # Callback to decode messages and dispatch to subscribers
+        self._dispatch = decoder_callback(self._decoder, self._subscribers, self._logger)
+
+    async def _do_hello(self, local_protocol_version: LocalProtocolVersion) -> LocalChannelParams | None:
+        """Perform the initial handshaking and return encoder params if successful."""
+        self._logger.debug(
+            "Attempting to use the %s protocol for client %s...",
+            local_protocol_version,
+            self._host,
+        )
+        request = RoborockMessage(
+            protocol=RoborockMessageProtocol.HELLO_REQUEST,
+            version=local_protocol_version.encode(),
+            random=self._params.connect_nonce,
+            seq=1,
+        )
+        try:
+            response = await self._send_message(
+                roborock_message=request,
+                request_id=request.seq,
+                response_protocol=RoborockMessageProtocol.HELLO_RESPONSE,
+            )
+            self._logger.debug(
+                "Client %s speaks the %s protocol.",
+                self._host,
+                local_protocol_version,
+            )
+            return LocalChannelParams(
+                local_key=self._params.local_key, connect_nonce=self._params.connect_nonce, ack_nonce=response.random
+            )
+        except RoborockException as e:
+            self._logger.debug(
+                "Client %s did not respond or does not speak the %s protocol. %s",
+                self._host,
+                local_protocol_version,
+                e,
+            )
+            return None
+
+    async def _hello(self):
+        """Send hello to the device to negotiate protocol."""
+        attempt_versions = [LocalProtocolVersion.V1, LocalProtocolVersion.L01]
+        if self._local_protocol_version:
+            # Sort to try the preferred version first
+            attempt_versions.sort(key=lambda v: v != self._local_protocol_version)
+
+        for version in attempt_versions:
+            params = await self._do_hello(version)
+            if params is not None:
+                self._local_protocol_version = version
+                self._update_encoder_decoder(params)
+                return
+
+        raise RoborockException("Failed to connect to device with any known protocol")
+
+    async def _ping(self) -> None:
+        ping_message = RoborockMessage(
+            protocol=RoborockMessageProtocol.PING_REQUEST, version=self.protocol_version.encode()
+        )
+        await self._send_message(
+            roborock_message=ping_message,
+            request_id=ping_message.seq,
+            response_protocol=RoborockMessageProtocol.PING_RESPONSE,
+        )
+
+    async def _keep_alive_loop(self) -> None:
+        while self._is_connected:
+            try:
+                await asyncio.sleep(_PING_INTERVAL)
+                if self._is_connected:
+                    await self._ping()
+            except asyncio.CancelledError:
+                break
+            except Exception:
+                self._logger.debug("Keep-alive ping failed", exc_info=True)
+                # Retry next interval
+
+    @property
+    def protocol_version(self) -> LocalProtocolVersion:
+        """Return the negotiated local protocol version, or a sensible default."""
+        if self._local_protocol_version is not None:
+            return self._local_protocol_version
+        return LocalProtocolVersion.V1
+
+    @property
+    def is_connected(self) -> bool:
+        """Check if the channel is currently connected."""
+        return self._is_connected
+
+    @property
+    def is_local_connected(self) -> bool:
+        """Check if the channel is currently connected locally."""
+        return self._is_connected
+
+    async def connect(self) -> None:
+        """Connect to the device and negotiate protocol."""
+        if self._is_connected:
+            self._logger.debug("Unexpected call to connect when already connected")
+            return
+        loop = get_running_loop()
+        protocol = _LocalProtocol(self._data_received, self._connection_lost)
+        try:
+            self._transport, self._protocol = await loop.create_connection(lambda: protocol, self._host, _PORT)
+            self._is_connected = True
+        except OSError as e:
+            raise RoborockConnectionException(f"Failed to connect to {self._host}:{_PORT}") from e
+
+        # Perform protocol negotiation
+        try:
+            await self._hello()
+            self._keep_alive_task = asyncio.create_task(self._keep_alive_loop())
+        except RoborockException:
+            # If protocol negotiation fails, clean up the connection state
+            self.close()
+            raise
+
+    def _data_received(self, data: bytes) -> None:
+        """Invoked when data is received on the stream."""
+        self._dispatch(data)
+
+    def close(self) -> None:
+        """Disconnect from the device."""
+        if self._keep_alive_task:
+            self._keep_alive_task.cancel()
+            self._keep_alive_task = None
+        if self._transport:
+            self._transport.close()
+        else:
+            self._logger.warning("Close called but transport is already None")
+        self._transport = None
+        self._is_connected = False
+
+    def _connection_lost(self, exc: Exception | None) -> None:
+        """Handle connection loss."""
+        self._logger.debug("Connection lost to %s", self._host, exc_info=exc)
+        if self._keep_alive_task:
+            self._keep_alive_task.cancel()
+            self._keep_alive_task = None
+        self._transport = None
+        self._is_connected = False
+
+    async def subscribe(self, callback: Callable[[RoborockMessage], None]) -> Callable[[], None]:
+        """Subscribe to all messages from the device."""
+        return self._subscribers.add_callback(callback)
+
+    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:
+            encoded_msg = self._encoder(message)
+        except Exception as err:
+            self._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)
+        except Exception as err:
+            self._logger.exception("Uncaught error sending command")
+            raise RoborockException(f"Failed to send message: {message}") from err
+
+    async def _send_message(
+        self,
+        roborock_message: RoborockMessage,
+        request_id: int,
+        response_protocol: int,
+    ) -> RoborockMessage:
+        """Send a raw message and wait for a raw response."""
+        future: asyncio.Future[RoborockMessage] = asyncio.Future()
+
+        def find_response(response_message: RoborockMessage) -> None:
+            if response_message.protocol == response_protocol and response_message.seq == request_id:
+                future.set_result(response_message)
+
+        unsub = await self.subscribe(find_response)
+        try:
+            await self.publish(roborock_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()
+
+
+# This module provides a factory function to create LocalChannel instances.
+#
+# TODO: Make a separate LocalSession and use it to manage retries with the host,
+# similar to how MqttSession works. For now this is a simple factory function
+# for creating channels.
+LocalSession = Callable[[str], LocalChannel]
+
+
+def create_local_session(local_key: str, device_uid: str) -> LocalSession:
+    """Creates a local session which can create local channels.
+
+    This plays a role similar to the MqttSession but is really just a factory
+    for creating LocalChannel instances with the same local key.
+    """
+
+    def create_local_channel(host: str) -> LocalChannel:
+        """Create a LocalChannel instance for the given host."""
+        return LocalChannel(host, local_key, device_uid)
+
+    return create_local_channel

roborock_cli/_vendor/roborock/devices/transport/mqtt_channel.py

@@ -0,0 +1,118 @@
+"""Modules for communicating with specific Roborock devices over MQTT."""
+
+import asyncio
+import logging
+from collections.abc import AsyncGenerator, Callable
+
+from roborock_cli._vendor.roborock.callbacks import decoder_callback
+from roborock_cli._vendor.roborock.data import HomeDataDevice, RRiot, UserData
+from roborock_cli._vendor.roborock.exceptions import RoborockException
+from roborock_cli._vendor.roborock.mqtt.health_manager import HealthManager
+from roborock_cli._vendor.roborock.mqtt.session import MqttParams, MqttSession, MqttSessionException
+from roborock_cli._vendor.roborock.protocol import create_mqtt_decoder, create_mqtt_encoder
+from roborock_cli._vendor.roborock.roborock_message import RoborockMessage
+from roborock_cli._vendor.roborock.util import RoborockLoggerAdapter
+
+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._logger = RoborockLoggerAdapter(duid=duid, logger=_LOGGER)
+        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 health_manager(self) -> HealthManager:
+        """Return the health manager for the session."""
+        return self._mqtt_session.health_manager
+
+    @property
+    def is_local_connected(self) -> bool:
+        """Return true if the channel is connected locally."""
+        return False
+
+    @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.
+        """
+        dispatch = decoder_callback(self._decoder, callback, _LOGGER)
+        return await self._mqtt_session.subscribe(self._subscribe_topic, dispatch)
+
+    async def subscribe_stream(self) -> AsyncGenerator[RoborockMessage, None]:
+        """Subscribe to the device's message stream.
+
+        This is useful for processing all incoming messages in an async for loop,
+        when they are not necessarily associated with a specific request.
+        """
+        message_queue: asyncio.Queue[RoborockMessage] = asyncio.Queue()
+        unsub = await self.subscribe(message_queue.put_nowait)
+        try:
+            while True:
+                message = await message_queue.get()
+                yield message
+        finally:
+            unsub()
+
+    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:
+            self._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:
+            self._logger.debug("Error publishing MQTT message: %s", e)
+            raise RoborockException(f"Failed to publish MQTT message: {e}") from e
+
+    async def restart(self) -> None:
+        """Restart the underlying MQTT session."""
+        await self._mqtt_session.restart()
+
+
+def create_mqtt_channel(
+    user_data: UserData, mqtt_params: MqttParams, mqtt_session: MqttSession, device: HomeDataDevice
+) -> MqttChannel:
+    """Create a MQTT channel for the given device."""
+    return MqttChannel(mqtt_session, device.duid, device.local_key, user_data.rriot, mqtt_params)

roborock_cli/_vendor/roborock/diagnostics.py

@@ -0,0 +1,166 @@
+"""Diagnostics for debugging.
+
+A Diagnostics object can be used to track counts and latencies of various
+operations within a module. This can be useful for debugging performance issues
+or understanding usage patterns.
+
+This is an internal facing module and is not intended for public use. Diagnostics
+data is collected and exposed to clients via higher level APIs like the
+DeviceManager.
+"""
+
+from __future__ import annotations
+
+import time
+from collections import Counter
+from collections.abc import Generator, Mapping
+from contextlib import contextmanager
+from typing import Any, TypeVar, cast
+
+
+class Diagnostics:
+    """A class that holds diagnostics information for a module.
+
+    You can use this class to hold counter or for recording timing information
+    that can be exported as a dictionary for debugging purposes.
+    """
+
+    def __init__(self) -> None:
+        """Initialize Diagnostics."""
+        self._counter: Counter = Counter()
+        self._subkeys: dict[str, Diagnostics] = {}
+
+    def increment(self, key: str, count: int = 1) -> None:
+        """Increment a counter for the specified key/event."""
+        self._counter.update(Counter({key: count}))
+
+    def elapsed(self, key_prefix: str, elapsed_ms: int = 1) -> None:
+        """Track a latency event for the specified key/event prefix."""
+        self.increment(f"{key_prefix}_count", 1)
+        self.increment(f"{key_prefix}_sum", elapsed_ms)
+
+    def as_dict(self) -> Mapping[str, Any]:
+        """Return diagnostics as a debug dictionary."""
+        data: dict[str, Any] = {k: self._counter[k] for k in self._counter}
+        for k, d in self._subkeys.items():
+            v = d.as_dict()
+            if not v:
+                continue
+            data[k] = v
+        return data
+
+    def subkey(self, key: str) -> Diagnostics:
+        """Return sub-Diagnostics object with the specified subkey.
+
+        This will create a new Diagnostics object if one does not already exist
+        for the specified subkey. Stats from the sub-Diagnostics will be included
+        in the parent Diagnostics when exported as a dictionary.
+
+        Args:
+            key: The subkey for the diagnostics.
+
+        Returns:
+            The Diagnostics object for the specified subkey.
+        """
+        if key not in self._subkeys:
+            self._subkeys[key] = Diagnostics()
+        return self._subkeys[key]
+
+    @contextmanager
+    def timer(self, key_prefix: str) -> Generator[None, None, None]:
+        """A context manager that records the timing of operations as a diagnostic."""
+        start = time.perf_counter()
+        try:
+            yield
+        finally:
+            end = time.perf_counter()
+            ms = int((end - start) * 1000)
+            self.elapsed(key_prefix, ms)
+
+    def reset(self) -> None:
+        """Clear all diagnostics, for testing."""
+        self._counter = Counter()
+        for d in self._subkeys.values():
+            d.reset()
+
+
+T = TypeVar("T")
+
+REDACT_KEYS = {
+    # Potential identifiers
+    "localKey",
+    "mac",
+    "bssid",
+    "sn",
+    "ip",
+    "u",
+    "s",
+    "h",
+    "k",
+    # Large binary blobs are entirely omitted
+    "imageContent",
+    "mapData",
+    "rawApiResponse",
+    # Home data
+    "id",  # We want to redact home_data.id but keep some other ids, see below
+    "name",
+    "productId",
+    "ipAddress",
+    "wifiName",
+    "lat",
+    "long",
+}
+KEEP_KEYS = {
+    # Product information not unique per user
+    "product.id",
+    "product.schema.id",
+    "product.schema.name",
+    # Room ids are likely unique per user, but don't seem too sensitive and are
+    # useful for debugging
+    "rooms.id",
+}
+DEVICE_UID = "duid"
+REDACTED = "**REDACTED**"
+
+
+def redact_device_data(data: T, path: str = "") -> T | dict[str, Any]:
+    """Redact sensitive data in a dict."""
+    if not isinstance(data, (Mapping, list)):
+        return data
+
+    if isinstance(data, list):
+        return cast(T, [redact_device_data(item, path) for item in data])
+
+    redacted = {**data}
+
+    for key, value in redacted.items():
+        curr_path = f"{path}.{key}" if path else key
+        if key in KEEP_KEYS or curr_path in KEEP_KEYS:
+            continue
+        if key in REDACT_KEYS or curr_path in REDACT_KEYS:
+            redacted[key] = REDACTED
+        elif key == DEVICE_UID and isinstance(value, str):
+            redacted[key] = redact_device_uid(value)
+        elif isinstance(value, dict):
+            redacted[key] = redact_device_data(value, curr_path)
+        elif isinstance(value, list):
+            redacted[key] = [redact_device_data(item, curr_path) for item in value]
+
+    return redacted
+
+
+def redact_topic_name(topic: str) -> str:
+    """Redact potentially identifying information from a topic name."""
+    parts = topic.split("/")
+    redacted_parts = parts[:4]
+    for part in parts[4:]:
+        if len(part) <= 5:
+            redacted_parts.append("*****")
+        else:
+            redacted_parts.append("*****" + part[-5:])
+    return "/".join(redacted_parts)
+
+
+def redact_device_uid(duid: str) -> str:
+    """Redact a device UID to hide identifying information."""
+    return "******" + duid[-5:]

roborock_cli/_vendor/roborock/exceptions.py

@@ -0,0 +1,95 @@
+"""Roborock exceptions."""
+
+from __future__ import annotations
+
+
+class RoborockException(Exception):
+    """Class for Roborock exceptions."""
+
+
+class RoborockTimeout(RoborockException):
+    """Class for Roborock timeout exceptions."""
+
+
+class RoborockConnectionException(RoborockException):
+    """Class for Roborock connection exceptions."""
+
+
+class RoborockBackoffException(RoborockException):
+    """Class for Roborock exceptions when many retries were made."""
+
+
+class VacuumError(RoborockException):
+    """Class for vacuum errors."""
+
+
+class CommandVacuumError(RoborockException):
+    """Class for command vacuum errors."""
+
+    def __init__(self, command: str | None, vacuum_error: VacuumError):
+        self.message = f"{command or 'unknown'}: {str(vacuum_error)}"
+        super().__init__(self.message)
+
+
+class UnknownMethodError(RoborockException):
+    """Class for an invalid method being sent."""
+
+
+class RoborockAccountDoesNotExist(RoborockException):
+    """Class for Roborock account does not exist exceptions."""
+
+
+class RoborockUrlException(RoborockException):
+    """Class for being unable to get the URL for the Roborock account."""
+
+
+class RoborockInvalidCode(RoborockException):
+    """Class for Roborock invalid code exceptions."""
+
+
+class RoborockInvalidEmail(RoborockException):
+    """Class for Roborock invalid formatted email exceptions."""
+
+
+class RoborockInvalidUserAgreement(RoborockException):
+    """Class for Roborock invalid user agreement exceptions."""
+
+
+class RoborockNoUserAgreement(RoborockException):
+    """Class for Roborock no user agreement exceptions."""
+
+
+class RoborockInvalidCredentials(RoborockException):
+    """Class for Roborock credentials have expired or changed."""
+
+
+class RoborockTooFrequentCodeRequests(RoborockException):
+    """Class for Roborock too frequent code requests exceptions."""
+
+
+class RoborockMissingParameters(RoborockException):
+    """Class for Roborock missing parameters exceptions."""
+
+
+class RoborockTooManyRequest(RoborockException):
+    """Class for Roborock too many request exceptions."""
+
+
+class RoborockRateLimit(RoborockException):
+    """Class for our rate limits exceptions."""
+
+
+class RoborockNoResponseFromBaseURL(RoborockException):
+    """We could not find an url that had a record of the given account."""
+
+
+class RoborockDeviceBusy(RoborockException):
+    """Class for Roborock device busy exceptions."""
+
+
+class RoborockInvalidStatus(RoborockException):
+    """Class for Roborock invalid status exceptions (device action locked)."""
+
+
+class RoborockUnsupportedFeature(RoborockException):
+    """Class for Roborock unsupported feature exceptions."""

roborock_cli/_vendor/roborock/map/__init__.py

@@ -0,0 +1,7 @@
+"""Module for Roborock map related data classes."""
+
+from .map_parser import MapParserConfig, ParsedMapData
+
+__all__ = [
+    "MapParserConfig",
+]