feat: add some basic B01 support (#429)

* feat: add some basic B01 support

* fix: lint

* fix: lint

---------

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

Author: Lash-L

roborock/devices/b01_channel.py

@@ -0,0 +1,30 @@
+"""Thin wrapper around the MQTT channel for Roborock B01 devices."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from roborock.protocols.b01_protocol import (
+    CommandType,
+    ParamsType,
+    decode_rpc_response,
+    encode_mqtt_payload,
+)
+
+from .mqtt_channel import MqttChannel
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def send_decoded_command(
+    mqtt_channel: MqttChannel,
+    dps: int,
+    command: CommandType,
+    params: ParamsType,
+) -> dict[int, 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)
+    response = await mqtt_channel.send_message(roborock_message)
+    return decode_rpc_response(response)  # type: ignore[return-value]

roborock/devices/device_manager.py

@@ -21,6 +21,7 @@
 from .cache import Cache, NoCache
 from .channel import Channel
 from .mqtt_channel import create_mqtt_channel
+from .traits.b01.props import B01PropsApi
 from .traits.dyad import DyadApi
 from .traits.status import StatusTrait
 from .traits.trait import Trait
@@ -45,6 +46,7 @@ class DeviceVersion(enum.StrEnum):
 
     V1 = "1.0"
     A01 = "A01"
+    B01 = "B01"
     UNKNOWN = "unknown"
 
 
@@ -159,6 +161,9 @@ def device_creator(device: HomeDataDevice, product: HomeDataProduct) -> Roborock
                         traits.append(ZeoApi(mqtt_channel))
                     case _:
                         raise NotImplementedError(f"Device {device.name} has unsupported category {product.category}")
+            case DeviceVersion.B01:
+                mqtt_channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
+                traits.append(B01PropsApi(mqtt_channel))
             case _:
                 raise NotImplementedError(f"Device {device.name} has unsupported version {device.pv}")
         return RoborockDevice(device, channel, traits)

roborock/devices/traits/b01/__init__.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from roborock import RoborockB01Methods
+from roborock.roborock_message import RoborockB01Props
+
+from ...b01_channel import send_decoded_command
+from ...mqtt_channel import MqttChannel
+from ..trait import Trait
+
+_LOGGER = logging.getLogger(__name__)
+
+__all__ = [
+    "B01PropsApi",
+]
+
+
+class B01PropsApi(Trait):
+    """API for interacting with B01 devices."""
+
+    name = "B01_props"
+
+    def __init__(self, channel: MqttChannel) -> None:
+        """Initialize the B01Props API."""
+        self._channel = channel
+
+    async def query_values(self, props: list[RoborockB01Props]) -> dict[int, Any]:
+        """Query the device for the values of the given Dyad protocols."""
+        return await send_decoded_command(self._channel, dps=10000, command=RoborockB01Methods.GET_PROP, params=props)

roborock/devices/traits/b01/props.py

@@ -39,6 +39,7 @@
 _LOGGER = logging.getLogger(__name__)
 SALT = b"TXdfu$jyZ#TZHsg4"
 A01_HASH = "726f626f726f636b2d67a6d6da"
+B01_HASH = "5wwh9ikChRjASpMU8cxg7o1d2E"
 BROADCAST_TOKEN = b"qWKYcdQWrbm9hPqe"
 AP_CONFIG = 1
 SOCK_DISCOVERY = 2
@@ -213,6 +214,10 @@ def _encode(self, obj, context, _):
             decipher = AES.new(bytes(context.search("local_key"), "utf-8"), AES.MODE_CBC, bytes(iv, "utf-8"))
             f = decipher.encrypt(obj)
             return f
+        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)
         token = self.token_func(context)
         encrypted = Utils.encrypt_ecb(obj, token)
         return encrypted
@@ -224,6 +229,10 @@ def _decode(self, obj, context, _):
             decipher = AES.new(bytes(context.search("local_key"), "utf-8"), AES.MODE_CBC, bytes(iv, "utf-8"))
             f = decipher.decrypt(obj)
             return f
+        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)
         token = self.token_func(context)
         decrypted = Utils.decrypt_ecb(obj, token)
         return decrypted
@@ -248,7 +257,7 @@ class PrefixedStruct(Struct):
     def _parse(self, stream, context, path):
         subcon1 = Peek(Optional(Bytes(3)))
         peek_version = subcon1.parse_stream(stream, **context)
-        if peek_version not in (b"1.0", b"A01"):
+        if peek_version not in (b"1.0", b"A01", b"B01"):
             subcon2 = Bytes(4)
             subcon2.parse_stream(stream, **context)
         return super()._parse(stream, context, path)

roborock/protocol.py

@@ -0,0 +1,55 @@
+"""Roborock B01 Protocol encoding and decoding."""
+
+import json
+import logging
+from typing import Any
+
+from Crypto.Cipher import AES
+from Crypto.Util.Padding import pad, unpad
+
+from roborock import RoborockB01Methods
+from roborock.exceptions import RoborockException
+from roborock.roborock_message import (
+    RoborockMessage,
+    RoborockMessageProtocol,
+)
+
+_LOGGER = logging.getLogger(__name__)
+
+B01_VERSION = b"B01"
+CommandType = RoborockB01Methods | str
+ParamsType = list | dict | int | None
+
+
+def encode_mqtt_payload(dps: int, command: CommandType, params: ParamsType) -> RoborockMessage:
+    """Encode payload for B01 commands over MQTT."""
+    dps_data = {"dps": {dps: {"method": command, "params": params or []}}}
+    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 as err:
+        raise RoborockException(f"Unable to unpad B01 payload: {err}")
+
+    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/protocols/b01_protocol.py

@@ -4,6 +4,7 @@
 import math
 import time
 from dataclasses import dataclass, field
+from enum import StrEnum
 
 from roborock import RoborockEnum
 from roborock.util import get_next_int
@@ -130,6 +131,98 @@ class RoborockZeoProtocol(RoborockEnum):
     RPC_RESp = 10102
 
 
+class RoborockB01Protocol(RoborockEnum):
+    RPC_REQUEST = 101
+    RPC_RESPONSE = 102
+    ERROR_CODE = 120
+    STATE = 121
+    BATTERY = 122
+    FAN_POWER = 123
+    WATER_BOX_MODE = 124
+    MAIN_BRUSH_LIFE = 125
+    SIDE_BRUSH_LIFE = 126
+    FILTER_LIFE = 127
+    OFFLINE_STATUS = 135
+    CLEAN_TIMES = 136
+    CLEANING_PREFERENCE = 137
+    CLEAN_TASK_TYPE = 138
+    BACK_TYPE = 139
+    DOCK_TASK_TYPE = 140
+    CLEANING_PROGRESS = 141
+    FC_STATE = 142
+    START_CLEAN_TASK = 201
+    START_BACK_DOCK_TASK = 202
+    START_DOCK_TASK = 203
+    PAUSE = 204
+    RESUME = 205
+    STOP = 206
+    CEIP = 207
+
+
+class RoborockB01Props(StrEnum):
+    """Properties requested by the Roborock B01 model."""
+
+    STATUS = "status"
+    FAULT = "fault"
+    WIND = "wind"
+    WATER = "water"
+    MODE = "mode"
+    QUANTITY = "quantity"
+    ALARM = "alarm"
+    VOLUME = "volume"
+    HYPA = "hypa"
+    MAIN_BRUSH = "main_brush"
+    SIDE_BRUSH = "side_brush"
+    MOP_LIFE = "mop_life"
+    MAIN_SENSOR = "main_sensor"
+    NET_STATUS = "net_status"
+    REPEAT_STATE = "repeat_state"
+    TANK_STATE = "tank_state"
+    SWEEP_TYPE = "sweep_type"
+    CLEAN_PATH_PREFERENCE = "clean_path_preference"
+    CLOTH_STATE = "cloth_state"
+    TIME_ZONE = "time_zone"
+    TIME_ZONE_INFO = "time_zone_info"
+    LANGUAGE = "language"
+    CLEANING_TIME = "cleaning_time"
+    REAL_CLEAN_TIME = "real_clean_time"
+    CLEANING_AREA = "cleaning_area"
+    CUSTOM_TYPE = "custom_type"
+    SOUND = "sound"
+    WORK_MODE = "work_mode"
+    STATION_ACT = "station_act"
+    CHARGE_STATE = "charge_state"
+    CURRENT_MAP_ID = "current_map_id"
+    MAP_NUM = "map_num"
+    DUST_ACTION = "dust_action"
+    QUIET_IS_OPEN = "quiet_is_open"
+    QUIET_BEGIN_TIME = "quiet_begin_time"
+    QUIET_END_TIME = "quiet_end_time"
+    CLEAN_FINISH = "clean_finish"
+    VOICE_TYPE = "voice_type"
+    VOICE_TYPE_VERSION = "voice_type_version"
+    ORDER_TOTAL = "order_total"
+    BUILD_MAP = "build_map"
+    PRIVACY = "privacy"
+    DUST_AUTO_STATE = "dust_auto_state"
+    DUST_FREQUENCY = "dust_frequency"
+    CHILD_LOCK = "child_lock"
+    MULTI_FLOOR = "multi_floor"
+    MAP_SAVE = "map_save"
+    LIGHT_MODE = "light_mode"
+    GREEN_LASER = "green_laser"
+    DUST_BAG_USED = "dust_bag_used"
+    ORDER_SAVE_MODE = "order_save_mode"
+    MANUFACTURER = "manufacturer"
+    BACK_TO_WASH = "back_to_wash"
+    CHARGE_STATION_TYPE = "charge_station_type"
+    PV_CUT_CHARGE = "pv_cut_charge"
+    PV_CHARGING = "pv_charging"
+    SERIAL_NUMBER = "serial_number"
+    RECOMMEND = "recommend"
+    ADD_SWEEP_STATUS = "add_sweep_status"
+
+
 ROBOROCK_DATA_STATUS_PROTOCOL = [
     RoborockDataProtocol.ERROR_CODE,
     RoborockDataProtocol.STATE,

roborock/roborock_message.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from enum import Enum
+from enum import Enum, StrEnum
 
 from .containers import (
     CleanRecord,
@@ -271,6 +271,20 @@ class RoborockCommand(str, Enum):
     APP_GET_ROBOT_SETTING = "app_get_robot_setting"
 
 
+class RoborockB01Methods(StrEnum):
+    """Methods used by the Roborock B01 model."""
+
+    GET_PROP = "prop.get"
+    GET_MAP_LIST = "service.get_map_list"
+    UPLOAD_BY_MAPTYPE = "service.upload_by_maptype"
+    SET_PROP = "prop.set"
+    GET_PREFERENCE = "service.get_preference"
+    GET_RECORD_LIST = "service.get_record_list"
+    GET_ORDER = "service.get_order"
+    EVENT_ORDER_LIST_POST = "event.order_list.post"
+    POST_PROP = "prop.post"
+
+
 @dataclass
 class DockSummary(RoborockBase):
     dust_collection_mode: DustCollectionMode | None = None