feat: add b01 Q7 basic getter support (#662)

* feat: add b01 Q7 support

* chore: small tweaks

* fix: test

* chore: tests

* chore: fix test naming

* chore: update snapshot

* chore: update roborock/devices/b01_channel.py

Co-authored-by: Allen Porter <allen.porter@gmail.com>

* chore: apply suggestions from code review

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Allen Porter <allen.porter@gmail.com>

* chore: address some comments

* chore: rename

---------

Co-authored-by: Allen Porter <allen.porter@gmail.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 3 people

roborock/data/b01_q7/b01_q7_containers.py

@@ -69,62 +69,62 @@ class B01Props(RoborockBase):
     This dataclass is generated based on the device's status JSON object.
     """
 
-    status: WorkStatusMapping
-    fault: B01Fault
-    wind: SCWindMapping
-    water: int
-    mode: int
-    quantity: int
-    alarm: int
-    volume: int
-    hypa: int
-    main_brush: int
-    side_brush: int
-    mop_life: int
-    main_sensor: int
-    net_status: NetStatus
-    repeat_state: int
-    tank_state: int
-    sweep_type: int
-    clean_path_preference: int
-    cloth_state: int
-    time_zone: int
-    time_zone_info: str
-    language: int
-    cleaning_time: int
-    real_clean_time: int
-    cleaning_area: int
-    custom_type: int
-    sound: int
-    work_mode: WorkModeMapping
-    station_act: int
-    charge_state: int
-    current_map_id: int
-    map_num: int
-    dust_action: int
-    quiet_is_open: int
-    quiet_begin_time: int
-    quiet_end_time: int
-    clean_finish: int
-    voice_type: int
-    voice_type_version: int
-    order_total: OrderTotal
-    build_map: int
-    privacy: Privacy
-    dust_auto_state: int
-    dust_frequency: int
-    child_lock: int
-    multi_floor: int
-    map_save: int
-    light_mode: int
-    green_laser: int
-    dust_bag_used: int
-    order_save_mode: int
-    manufacturer: str
-    back_to_wash: int
-    charge_station_type: int
-    pv_cut_charge: int
-    pv_charging: PvCharging
-    serial_number: str
-    recommend: Recommend
-    add_sweep_status: int
+    status: WorkStatusMapping | None = None
+    fault: B01Fault | None = None
+    wind: SCWindMapping | None = None
+    water: int | None = None
+    mode: int | None = None
+    quantity: int | None = None
+    alarm: int | None = None
+    volume: int | None = None
+    hypa: int | None = None
+    main_brush: int | None = None
+    side_brush: int | None = None
+    mop_life: int | None = None
+    main_sensor: int | None = None
+    net_status: NetStatus | None = None
+    repeat_state: int | None = None
+    tank_state: int | None = None
+    sweep_type: int | None = None
+    clean_path_preference: int | None = None
+    cloth_state: int | None = None
+    time_zone: int | None = None
+    time_zone_info: str | None = None
+    language: int | None = None
+    cleaning_time: int | None = None
+    real_clean_time: int | None = None
+    cleaning_area: int | None = None
+    custom_type: int | None = None
+    sound: int | None = None
+    work_mode: WorkModeMapping | None = None
+    station_act: int | None = None
+    charge_state: int | None = None
+    current_map_id: int | None = None
+    map_num: int | None = None
+    dust_action: int | None = None
+    quiet_is_open: int | None = None
+    quiet_begin_time: int | None = None
+    quiet_end_time: int | None = None
+    clean_finish: int | None = None
+    voice_type: int | None = None
+    voice_type_version: int | None = None
+    order_total: OrderTotal | None = None
+    build_map: int | None = None
+    privacy: Privacy | None = None
+    dust_auto_state: int | None = None
+    dust_frequency: int | None = None
+    child_lock: int | None = None
+    multi_floor: int | None = None
+    map_save: int | None = None
+    light_mode: int | None = None
+    green_laser: int | None = None
+    dust_bag_used: int | None = None
+    order_save_mode: int | None = None
+    manufacturer: str | None = None
+    back_to_wash: int | None = None
+    charge_station_type: int | None = None
+    pv_cut_charge: int | None = None
+    pv_charging: PvCharging | None = None
+    serial_number: str | None = None
+    recommend: Recommend | None = None
+    add_sweep_status: int | None = None

roborock/devices/b01_channel.py

@@ -2,26 +2,76 @@
 
 from __future__ import annotations
 
+import asyncio
+import json
 import logging
+from typing import Any
 
+from roborock.exceptions import RoborockException
 from roborock.protocols.b01_protocol import (
     CommandType,
     ParamsType,
+    decode_rpc_response,
     encode_mqtt_payload,
 )
+from roborock.roborock_message import RoborockMessage
+from roborock.util import get_next_int
 
 from .mqtt_channel import MqttChannel
 
 _LOGGER = logging.getLogger(__name__)
+_TIMEOUT = 10.0
 
 
 async def send_decoded_command(
     mqtt_channel: MqttChannel,
     dps: int,
     command: CommandType,
     params: ParamsType,
-) -> None:
+) -> dict[str, 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(dps, command, params)
-    await mqtt_channel.publish(roborock_message)
+    msg_id = str(get_next_int(100000000000, 999999999999))
+    roborock_message = encode_mqtt_payload(dps, command, params, msg_id)
+    future: asyncio.Future[dict[str, 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.info("Failed to decode b01 message: %s: %s", 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") == msg_id:
+                _LOGGER.debug("Received query response: %s", inner)
+                data = inner.get("data")
+                if not future.done():
+                    if isinstance(data, dict):
+                        future.set_result(data)
+                    else:
+                        future.set_exception(RoborockException(f"Unexpected data type for response: {data}"))
+
+    unsub = await mqtt_channel.subscribe(find_response)
+
+    _LOGGER.debug("Sending MQTT message: %s", roborock_message)
+    try:
+        await mqtt_channel.publish(roborock_message)
+        try:
+            return await asyncio.wait_for(future, timeout=_TIMEOUT)
+        except TimeoutError as ex:
+            raise RoborockException(f"Command timed out after {_TIMEOUT}s") from ex
+    finally:
+        unsub()

roborock/devices/traits/b01/q7/__init__.py

@@ -1,6 +1,7 @@
 """Traits for Q7 B01 devices.
 Potentially other devices may fall into this category in the future."""
 
+from roborock import B01Props
 from roborock.devices.b01_channel import send_decoded_command
 from roborock.devices.mqtt_channel import MqttChannel
 from roborock.devices.traits import Trait
@@ -13,17 +14,18 @@
 
 
 class Q7PropertiesApi(Trait):
-    """API for interacting with Q7 B01 devices."""
+    """API for interacting with B01 devices."""
 
     def __init__(self, channel: MqttChannel) -> None:
         """Initialize the B01Props API."""
         self._channel = channel
 
-    async def query_values(self, props: list[RoborockB01Props]) -> None:
+    async def query_values(self, props: list[RoborockB01Props]) -> B01Props | None:
         """Query the device for the values of the given Q7 properties."""
-        await send_decoded_command(
+        result = await send_decoded_command(
             self._channel, dps=10000, command=RoborockB01Q7Methods.GET_PROP, params={"property": props}
         )
+        return B01Props.from_dict(result)
 
 
 def create(channel: MqttChannel) -> Q7PropertiesApi:

roborock/protocol.py

@@ -276,12 +276,11 @@ def _encode(self, obj, context, _):
         if context.version == b"A01":
             iv = md5hex(format(context.random, "08x") + A01_HASH)[8:24]
             decipher = AES.new(bytes(context.search("local_key"), "utf-8"), AES.MODE_CBC, bytes(iv, "utf-8"))
-            f = decipher.encrypt(obj)
-            return f
+            return decipher.encrypt(pad(obj, AES.block_size))
         elif context.version == b"B01":
             iv = md5hex(f"{context.random:08x}" + B01_HASH)[9:25]
             decipher = AES.new(bytes(context.search("local_key"), "utf-8"), AES.MODE_CBC, bytes(iv, "utf-8"))
-            return decipher.encrypt(obj)
+            return decipher.encrypt(pad(obj, AES.block_size))
         elif context.version == b"L01":
             return Utils.encrypt_gcm_l01(
                 plaintext=obj,
@@ -301,12 +300,11 @@ def _decode(self, obj, context, _):
         if context.version == b"A01":
             iv = md5hex(format(context.random, "08x") + A01_HASH)[8:24]
             decipher = AES.new(bytes(context.search("local_key"), "utf-8"), AES.MODE_CBC, bytes(iv, "utf-8"))
-            f = decipher.decrypt(obj)
-            return f
+            return unpad(decipher.decrypt(obj), AES.block_size)
         elif context.version == b"B01":
             iv = md5hex(f"{context.random:08x}" + B01_HASH)[9:25]
             decipher = AES.new(bytes(context.search("local_key"), "utf-8"), AES.MODE_CBC, bytes(iv, "utf-8"))
-            return decipher.decrypt(obj)
+            return unpad(decipher.decrypt(obj), AES.block_size)
         elif context.version == b"L01":
             return Utils.decrypt_gcm_l01(
                 payload=obj,
@@ -350,6 +348,11 @@ def _parse(self, stream, context, path):
             # Read remaining data to find a valid header
             data = stream.read()
 
+            if not data:
+                # EOF reached, let the parser fail naturally without logging
+                stream_seek(stream, current_pos, 0, path)
+                return super()._parse(stream, context, path)
+
             start_index = -1
             # Find the earliest occurrence of any valid version in a single pass
             for i in range(len(data) - 2):

roborock/protocols/b01_protocol.py

@@ -13,7 +13,6 @@
     RoborockMessage,
     RoborockMessageProtocol,
 )
-from roborock.util import get_next_int
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -22,13 +21,13 @@
 ParamsType = list | dict | int | None
 
 
-def encode_mqtt_payload(dps: int, command: CommandType, params: ParamsType) -> RoborockMessage:
+def encode_mqtt_payload(dps: int, command: CommandType, params: ParamsType, msg_id: str) -> RoborockMessage:
     """Encode payload for B01 commands over MQTT."""
     dps_data = {
         "dps": {
             dps: {
                 "method": str(command),
-                "msgId": str(get_next_int(100000000000, 999999999999)),
+                "msgId": msg_id,
                 "params": params or [],
             }
         }
@@ -47,8 +46,9 @@ def decode_rpc_response(message: RoborockMessage) -> dict[int, Any]:
         raise RoborockException("Invalid B01 message format: missing payload")
     try:
         unpadded = unpad(message.payload, AES.block_size)
-    except ValueError as err:
-        raise RoborockException(f"Unable to unpad B01 payload: {err}")
+    except ValueError:
+        # It would be better to fail down the line.
+        unpadded = message.payload
 
     try:
         payload = json.loads(unpadded.decode())