feat: Add a CLI for exercising the asyncio MQTT session

Author: allenporter

roborock/cli.py

@@ -4,6 +4,7 @@
 import logging
 from pathlib import Path
 from typing import Any
+import asyncio
 
 import click
 from pyshark import FileCapture  # type: ignore
@@ -12,8 +13,10 @@
 
 from roborock import RoborockException
 from roborock.containers import DeviceData, HomeDataProduct, LoginData
-from roborock.protocol import MessageParser
+from roborock.protocol import MessageParser, create_mqtt_params
 from roborock.util import run_sync
+from roborock.mqtt.session import MqttParams, MqttSession
+from roborock.mqtt.roborock_session import create_mqtt_session
 from roborock.version_1_apis.roborock_local_client_v1 import RoborockLocalClientV1
 from roborock.version_1_apis.roborock_mqtt_client_v1 import RoborockMqttClientV1
 from roborock.web_api import RoborockApiClient
@@ -45,7 +48,8 @@ def validate(self):
         if self._login_data is None:
             raise RoborockException("You must login first")
 
-    def login_data(self):
+    def login_data(self) -> LoginData:
+        """Get the login data."""
         self.validate()
         return self._login_data
 
@@ -79,6 +83,54 @@ async def login(ctx, email, password):
     context.update(LoginData(user_data=user_data, email=email))
 
 
+@click.command()
+@click.pass_context
+@click.option("--duration", default=10, help="Duration to run the MQTT session in seconds")
+@run_sync()
+async def session(ctx, duration: int):
+    context: RoborockContext = ctx.obj
+    login_data = context.login_data()
+
+    # Discovery devices if not already available
+    if not login_data.home_data:
+        await _discover(ctx)
+        login_data = context.login_data()
+    if not login_data.home_data or not login_data.home_data.devices:
+        raise RoborockException("Unable to discover devices")
+
+    all_devices = login_data.home_data.devices + login_data.home_data.received_devices
+    click.echo(f"Discovered devices: {', '.join([device.name for device in all_devices])}")
+
+    rriot = login_data.user_data.rriot
+    params = create_mqtt_params(rriot)
+
+    mqtt_session = await create_mqtt_session(params)
+    click.echo("Starting MQTT session...")
+    if not mqtt_session.connected:
+        raise RoborockException("Failed to connect to MQTT broker")
+
+    def on_message(bytes: bytes):
+        """Callback function to handle incoming MQTT messages."""
+        # Decode the first 20 bytes of the message for display
+        bytes = bytes[:20]
+        click.echo(f"Received message: b\"{bytes.decode('utf-8')}...\"")
+
+    unsubs = []
+    for device in all_devices:
+        device_topic = f"rr/m/o/{rriot.u}/{params.username}/{device.duid}"
+        unsub = await mqtt_session.subscribe(device_topic, on_message)
+        unsubs.append(unsub)
+
+    click.echo("MQTT session started. Listening for messages...")
+    await asyncio.sleep(duration)
+
+    click.echo("Stopping MQTT session...")
+    for unsub in unsubs:
+        unsub()
+    await mqtt_session.close()
+
+
+
 async def _discover(ctx):
     context: RoborockContext = ctx.obj
     login_data = context.login_data()
@@ -253,6 +305,7 @@ def on_package(packet: Packet):
 cli.add_command(status)
 cli.add_command(command)
 cli.add_command(parser)
+cli.add_command(session)
 
 
 def main():

roborock/containers.py

@@ -9,6 +9,7 @@
 from enum import Enum
 from typing import Any, NamedTuple, get_args, get_origin
 
+from .mqtt.session import MqttParams
 from .code_mappings import (
     RoborockCategory,
     RoborockCleanType,

roborock/mqtt/roborock_session.py

@@ -116,9 +116,9 @@ async def _run_task(self, start_future: asyncio.Future[None] | None) -> None:
                 _LOGGER.info("MQTT error: %s", err)
             except asyncio.CancelledError as err:
                 if start_future:
-                    _LOGGER.debug("MQTT loop was cancelled")
+                    _LOGGER.debug("MQTT loop was cancelled while starting")
                     start_future.set_exception(err)
-                _LOGGER.debug("MQTT loop was cancelled whiel starting")
+                _LOGGER.debug("MQTT loop was cancelled")
                 return
             # Catch exceptions to avoid crashing the loop
             # and to allow the loop to retry.
@@ -171,8 +171,7 @@ async def _mqtt_client(self, params: MqttParams) -> aiomqtt.Client:
                 self._client = None
 
     async def _process_message_loop(self, client: aiomqtt.Client) -> None:
-        _LOGGER.debug("client=%s", client)
-        _LOGGER.debug("Processing MQTT messages: %s", client.messages)
+        _LOGGER.debug("Processing MQTT messages")
         async for message in client.messages:
             _LOGGER.debug("Received message: %s", message)
             for listener in self._listeners.get(message.topic.value, []):

roborock/protocol.py

@@ -8,6 +8,7 @@
 import logging
 from asyncio import BaseTransport, Lock
 from collections.abc import Callable
+from urllib.parse import urlparse
 
 from construct import (  # type: ignore
     Bytes,
@@ -30,8 +31,10 @@
 from Crypto.Cipher import AES
 from Crypto.Util.Padding import pad, unpad
 
-from roborock import BroadcastMessage, RoborockException
+from roborock.containers import BroadcastMessage, RRiot
+from roborock.exceptions import RoborockException
 from roborock.roborock_message import RoborockMessage
+from roborock.mqtt.session import MqttParams
 
 _LOGGER = logging.getLogger(__name__)
 SALT = b"TXdfu$jyZ#TZHsg4"
@@ -359,3 +362,20 @@ def build(
 
 MessageParser: _Parser = _Parser(_Messages, True)
 BroadcastParser: _Parser = _Parser(_BroadcastMessage, False)
+
+
+
+def create_mqtt_params(rriot: RRiot) -> MqttParams:
+    """Return the MQTT parameters for this user."""
+    url = urlparse(rriot.r.m)
+    if not isinstance(url.hostname, str):
+        raise RoborockException("Url parsing returned an invalid hostname")
+    hashed_user =  md5hex(rriot.u + ":" + rriot.k)[2:10]
+    hashed_password = md5hex(rriot.s + ":" + rriot.k)[16:]
+    return MqttParams(
+        host=str(url.hostname),
+        port=url.port,
+        tls=(url.scheme == "ssl"),
+        username=hashed_user,
+        password=hashed_password,
+    )