From a53e73ceef1330a13e5285030614afcc26dfc64b Mon Sep 17 00:00:00 2001 From: Luke Date: Wed, 9 Apr 2025 10:03:23 -0400 Subject: [PATCH 1/8] chore: init commit --- roborock/mqtt_manager.py | 0 roborock/roborock_device.py | 0 2 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 roborock/mqtt_manager.py create mode 100644 roborock/roborock_device.py diff --git a/roborock/mqtt_manager.py b/roborock/mqtt_manager.py new file mode 100644 index 00000000..e69de29b diff --git a/roborock/roborock_device.py b/roborock/roborock_device.py new file mode 100644 index 00000000..e69de29b From 10d6531ec24c18598daeacca13abedf61d506fb4 Mon Sep 17 00:00:00 2001 From: Luke Date: Wed, 9 Apr 2025 10:05:17 -0400 Subject: [PATCH 2/8] chore: init commit --- roborock/mqtt_manager.py | 113 ++++++++++++++++++++++++++++++++++ roborock/roborock_device.py | 117 ++++++++++++++++++++++++++++++++++++ 2 files changed, 230 insertions(+) diff --git a/roborock/mqtt_manager.py b/roborock/mqtt_manager.py index e69de29b..c134c26f 100644 --- a/roborock/mqtt_manager.py +++ b/roborock/mqtt_manager.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import asyncio +import dataclasses +import logging +from collections.abc import Coroutine +from typing import Callable, Self +from urllib.parse import urlparse + +import aiomqtt +from aiomqtt import TLSParameters + +from roborock import RoborockException, UserData +from roborock.protocol import MessageParser, md5hex + +from .containers import DeviceData + +LOGGER = logging.getLogger(__name__) + + +@dataclasses.dataclass +class ClientWrapper: + publish_function: Coroutine[None] + unsubscribe_function: Coroutine[None] + subscribe_function: Coroutine[None] + + +class RoborockMqttManager: + client_wrappers: dict[str, ClientWrapper] = {} + _instance: Self = None + + def __new__(cls) -> RoborockMqttManager: + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + async def connect(self, user_data: UserData): + # Add some kind of lock so we don't try to connect if we are already trying to connect the same account. + if user_data.rriot.u not in self.client_wrappers: + loop = asyncio.get_event_loop() + loop.create_task(self._new_connect(user_data)) + + async def _new_connect(self, user_data: UserData): + rriot = user_data.rriot + mqtt_user = rriot.u + hashed_user = md5hex(mqtt_user + ":" + rriot.k)[2:10] + url = urlparse(rriot.r.m) + if not isinstance(url.hostname, str): + raise RoborockException("Url parsing returned an invalid hostname") + mqtt_host = str(url.hostname) + mqtt_port = url.port + + mqtt_password = rriot.s + hashed_password = md5hex(mqtt_password + ":" + rriot.k)[16:] + LOGGER.debug("Connecting to %s for %s", mqtt_host, mqtt_user) + + async with aiomqtt.Client( + hostname=mqtt_host, + port=mqtt_port, + username=hashed_user, + password=hashed_password, + keepalive=60, + tls_params=TLSParameters(), + ) as client: + # TODO: Handle logic for when client loses connection + LOGGER.info("Connected to %s for %s", mqtt_host, mqtt_user) + callbacks: dict[str, Callable] = {} + device_map = {} + + async def publish(device: DeviceData, payload: bytes): + await client.publish(f"rr/m/i/{mqtt_user}/{hashed_user}/{device.device.duid}", payload=payload) + + async def subscribe(device: DeviceData, callback): + LOGGER.debug(f"Subscribing to rr/m/o/{mqtt_user}/{hashed_user}/{device.device.duid}") + await client.subscribe(f"rr/m/o/{mqtt_user}/{hashed_user}/{device.device.duid}") + LOGGER.debug(f"Subscribed to rr/m/o/{mqtt_user}/{hashed_user}/{device.device.duid}") + callbacks[device.device.duid] = callback + device_map[device.device.duid] = device + return + + async def unsubscribe(device: DeviceData): + await client.unsubscribe(f"rr/m/o/{mqtt_user}/{hashed_user}/{device.device.duid}") + + self.client_wrappers[user_data.rriot.u] = ClientWrapper( + publish_function=publish, unsubscribe_function=unsubscribe, subscribe_function=subscribe + ) + async for message in client.messages: + try: + device_id = message.topic.value.split("/")[-1] + device = device_map[device_id] + message = MessageParser.parse(message.payload, device.device.local_key) + callbacks[device_id](message) + except Exception: + ... + + async def disconnect(self, user_data: UserData): + await self.client_wrappers[user_data.rriot.u].disconnect() + + async def subscribe(self, user_data: UserData, device: DeviceData, callback): + if user_data.rriot.u not in self.client_wrappers: + await self.connect(user_data) + # add some kind of lock to make sure we don't subscribe until the connection is successful + await asyncio.sleep(2) + await self.client_wrappers[user_data.rriot.u].subscribe_function(device, callback) + + async def unsubscribe(self): + pass + + async def publish(self, user_data: UserData, device, payload: bytes): + LOGGER.debug("Publishing topic for %s, Message: %s", device.device.duid, payload) + if user_data.rriot.u not in self.client_wrappers: + await self.connect(user_data) + await self.client_wrappers[user_data.rriot.u].publish_function(device, payload) diff --git a/roborock/roborock_device.py b/roborock/roborock_device.py index e69de29b..c54f9ed3 100644 --- a/roborock/roborock_device.py +++ b/roborock/roborock_device.py @@ -0,0 +1,117 @@ +import base64 +import json +import logging +import math +import secrets +import time + +from . import RoborockCommand +from .containers import DeviceData, UserData +from .mqtt_manager import RoborockMqttManager +from .protocol import MessageParser, Utils +from .roborock_message import RoborockMessage, RoborockMessageProtocol +from .util import RoborockLoggerAdapter, get_next_int + +_LOGGER = logging.getLogger(__name__) + + +class RoborockDevice: + def __init__(self, user_data: UserData, device_info: DeviceData): + self.user_data = user_data + self.device_info = device_info + self.data = None + self._logger = RoborockLoggerAdapter(device_info.device.name, _LOGGER) + self._mqtt_endpoint = base64.b64encode(Utils.md5(user_data.rriot.k.encode())[8:14]).decode() + self._local_endpoint = "abc" + self._nonce = secrets.token_bytes(16) + self.manager = RoborockMqttManager() + self.update_commands = self.determine_supported_commands() + + def determine_supported_commands(self): + # All devices support these + supported_commands = { + RoborockCommand.GET_CONSUMABLE, + RoborockCommand.GET_STATUS, + RoborockCommand.GET_CLEAN_SUMMARY, + } + # Get what features we can from the feature_set info. + + # If a command is not described in feature_set, we should just add it anyways and then let it fail on the first call and remove it. + robot_new_features = int(self.device_info.device.feature_set) + new_feature_info_str = self.device_info.device.new_feature_set + if 33554432 & int(robot_new_features): + supported_commands.add(RoborockCommand.GET_DUST_COLLECTION_MODE) + if 2 & int(new_feature_info_str[-8:], 16): + # TODO: May not be needed as i think this can just be found in Status, but just POC + supported_commands.add(RoborockCommand.APP_GET_CLEAN_ESTIMATE_INFO) + return supported_commands + + async def connect(self): + """Connect via MQTT and Local if possible.""" + await self.manager.subscribe(self.user_data, self.device_info, self.on_message) + await self.update() + + async def update(self): + for cmd in self.update_commands: + await self.send_message(method=cmd) + + def _get_payload( + self, + method: RoborockCommand | str, + params: list | dict | int | None = None, + secured=False, + use_cloud: bool = False, + ): + timestamp = math.floor(time.time()) + request_id = get_next_int(10000, 32767) + inner = { + "id": request_id, + "method": method, + "params": params or [], + } + if secured: + inner["security"] = { + "endpoint": self._mqtt_endpoint if use_cloud else self._local_endpoint, + "nonce": self._nonce.hex().lower(), + } + payload = bytes( + json.dumps( + { + "dps": {"101": json.dumps(inner, separators=(",", ":"))}, + "t": timestamp, + }, + separators=(",", ":"), + ).encode() + ) + return request_id, timestamp, payload + + async def send_message( + self, method: RoborockCommand | str, params: list | dict | int | None = None, use_cloud: bool = True + ): + request_id, timestamp, payload = self._get_payload(method, params, True, use_cloud) + request_protocol = RoborockMessageProtocol.RPC_REQUEST + roborock_message = RoborockMessage(timestamp=timestamp, protocol=request_protocol, payload=payload) + + local_key = self.device_info.device.local_key + msg = MessageParser.build(roborock_message, local_key, False) + if use_cloud: + await self.manager.publish(self.user_data, self.device_info, msg) + else: + # Handle doing local commands + pass + + def on_message(self, message: RoborockMessage): + # If message is command not supported - remove from self.update_commands + + # If message is an error - log it? + + # If message is 'ok' - ignore it + + # If message is anything else - store ids, and map back to id to determine message type. + # Then update self.data + + # If we haven't received a message in X seconds, the device is likely offline. I think we can continue the connection, + # but we should have some way to mark ourselves as unavailable. + + # This should also probably be split with on_cloud_message and on_local_message. + print(message) From 9b41365194d6dbbb4a8a0d2bb03d3eb4eb4ff75e Mon Sep 17 00:00:00 2001 From: Luke Date: Thu, 10 Apr 2025 09:46:34 -0400 Subject: [PATCH 3/8] feat: add device_features to automatically determine what is supported. --- roborock/code_mappings.py | 140 ++++++++++- roborock/containers.py | 474 +++++++++++++++++++++++++++----------- 2 files changed, 480 insertions(+), 134 deletions(-) diff --git a/roborock/code_mappings.py b/roborock/code_mappings.py index 9f0736e8..140bf0d8 100644 --- a/roborock/code_mappings.py +++ b/roborock/code_mappings.py @@ -1,7 +1,7 @@ from __future__ import annotations import logging -from enum import Enum, IntEnum +from enum import Enum, IntEnum, StrEnum _LOGGER = logging.getLogger(__name__) completed_warnings = set() @@ -50,6 +50,144 @@ def items(cls: type[RoborockEnum]): return cls.as_dict().items() +class RoborockProductNickname(StrEnum): + """Enumeration of product nicknames.""" + + CORAL = "Coral" + CORALPRO = "CoralPro" + PEARL = "Pearl" + PEARLC = "PearlC" + PEARLE = "PearlE" + PEARLELITE = "PearlELite" + PEARLPLUS = "PearlPlus" + PEARLPLUSS = "PearlPlusS" + PEARLS = "PearlS" + PEARLSLITE = "PearlSLite" + RUBYPLUS = "RubyPlus" + RUBYSC = "RubySC" + RUBYSE = "RubySE" + RUBYSLITE = "RubySLite" + TANOS = "Tanos" + TANOSE = "TanosE" + TANOSS = "TanosS" + TANOSSC = "TanosSC" + TANOSSE = "TanosSE" + TANOSSMAX = "TanosSMax" + TANOSSLITE = "TanosSLite" + TANOSSPLUS = "TanosSPlus" + TANOSV = "TanosV" + TOPAZS = "TopazS" + TOPAZSC = "TopazSC" + TOPAZSPLUS = "TopazSPlus" + TOPAZSPOWER = "TopazSPower" + TOPAZSV = "TopazSV" + ULTRON = "Ultron" + ULTRONE = "UltronE" + ULTRONLITE = "UltronLite" + ULTRONSC = "UltronSC" + ULTRONSE = "UltronSE" + ULTRONSPLUS = "UltronSPlus" + ULTRONSV = "UltronSV" + VERDELITE = "Verdelite" + VIVIAN = "Vivian" + VIVIANC = "VivianC" + + +short_model_to_enum = { + # Pearl Series + "a103": RoborockProductNickname.PEARLC, + "a104": RoborockProductNickname.PEARLC, + "a116": RoborockProductNickname.PEARLPLUSS, + "a117": RoborockProductNickname.PEARLPLUSS, + "a136": RoborockProductNickname.PEARLPLUSS, + "a122": RoborockProductNickname.PEARLSLITE, + "a123": RoborockProductNickname.PEARLSLITE, + "a167": RoborockProductNickname.PEARLE, + "a168": RoborockProductNickname.PEARLE, + "a169": RoborockProductNickname.PEARLELITE, + "a170": RoborockProductNickname.PEARLELITE, + "a74": RoborockProductNickname.PEARL, + "a75": RoborockProductNickname.PEARL, + "a100": RoborockProductNickname.PEARLS, + "a101": RoborockProductNickname.PEARLS, + "a86": RoborockProductNickname.PEARLPLUS, + "a87": RoborockProductNickname.PEARLPLUS, + # Vivian Series + "a158": RoborockProductNickname.VIVIANC, + "a159": RoborockProductNickname.VIVIANC, + "a134": RoborockProductNickname.VIVIAN, + "a135": RoborockProductNickname.VIVIAN, + "a155": RoborockProductNickname.VIVIAN, + "a156": RoborockProductNickname.VIVIAN, + # Coral Series + "a143": RoborockProductNickname.CORALPRO, + "a144": RoborockProductNickname.CORALPRO, + "a20": RoborockProductNickname.CORAL, + "a21": RoborockProductNickname.CORAL, + # Ultron Series + "a73": RoborockProductNickname.ULTRONLITE, + "a85": RoborockProductNickname.ULTRONLITE, + "a94": RoborockProductNickname.ULTRONSC, + "a95": RoborockProductNickname.ULTRONSC, + "a124": RoborockProductNickname.ULTRONSE, + "a125": RoborockProductNickname.ULTRONSE, + "a139": RoborockProductNickname.ULTRONSE, + "a140": RoborockProductNickname.ULTRONSE, + "a68": RoborockProductNickname.ULTRONSPLUS, + "a69": RoborockProductNickname.ULTRONSPLUS, + "a70": RoborockProductNickname.ULTRONSPLUS, + "a50": RoborockProductNickname.ULTRON, + "a51": RoborockProductNickname.ULTRON, + "a72": RoborockProductNickname.ULTRONE, + "a84": RoborockProductNickname.ULTRONE, + "a96": RoborockProductNickname.ULTRONSV, + "a97": RoborockProductNickname.ULTRONSV, + # Verdelite Series + "a146": RoborockProductNickname.VERDELITE, + "a147": RoborockProductNickname.VERDELITE, + # Topaz Series + "a29": RoborockProductNickname.TOPAZS, + "a30": RoborockProductNickname.TOPAZS, + "a76": RoborockProductNickname.TOPAZS, + "a46": RoborockProductNickname.TOPAZSPLUS, + "a47": RoborockProductNickname.TOPAZSPLUS, + "a66": RoborockProductNickname.TOPAZSPLUS, + "a64": RoborockProductNickname.TOPAZSC, + "a65": RoborockProductNickname.TOPAZSC, + "a26": RoborockProductNickname.TOPAZSV, + "a27": RoborockProductNickname.TOPAZSV, + "a62": RoborockProductNickname.TOPAZSPOWER, + # Tanos Series + "a23": RoborockProductNickname.TANOSSPLUS, + "a24": RoborockProductNickname.TANOSSPLUS, + "a37": RoborockProductNickname.TANOSSLITE, + "a38": RoborockProductNickname.TANOSSLITE, + "a39": RoborockProductNickname.TANOSSC, + "a40": RoborockProductNickname.TANOSSC, + "a33": RoborockProductNickname.TANOSSE, + "a34": RoborockProductNickname.TANOSSE, + "a52": RoborockProductNickname.TANOSSMAX, + "t6": RoborockProductNickname.TANOS, + "s6": RoborockProductNickname.TANOS, + "t7": RoborockProductNickname.TANOSE, + "a11": RoborockProductNickname.TANOSE, + "t7p": RoborockProductNickname.TANOSV, + "a09": RoborockProductNickname.TANOSV, + "a10": RoborockProductNickname.TANOSV, + "a14": RoborockProductNickname.TANOSS, + "a15": RoborockProductNickname.TANOSS, + # Ruby Series + "t4": RoborockProductNickname.RUBYPLUS, + "s4": RoborockProductNickname.RUBYPLUS, + "p5": RoborockProductNickname.RUBYSC, + "a08": RoborockProductNickname.RUBYSC, + "a19": RoborockProductNickname.RUBYSE, + "p6": RoborockProductNickname.RUBYSLITE, + "s5e": RoborockProductNickname.RUBYSLITE, + "a05": RoborockProductNickname.RUBYSLITE, +} + + class RoborockStateCode(RoborockEnum): unknown = 0 starting = 1 diff --git a/roborock/containers.py b/roborock/containers.py index afbe14b3..d0be532f 100644 --- a/roborock/containers.py +++ b/roborock/containers.py @@ -6,7 +6,7 @@ import re from dataclasses import asdict, dataclass, field from datetime import timezone -from enum import Enum +from enum import Enum, IntEnum from typing import Any, NamedTuple, get_args, get_origin from .code_mappings import ( @@ -43,8 +43,10 @@ RoborockMopModeS7, RoborockMopModeS8MaxVUltra, RoborockMopModeS8ProUltra, + RoborockProductNickname, RoborockStartType, RoborockStateCode, + short_model_to_enum, ) from .const import ( CLEANING_BRUSH_REPLACE_TIME, @@ -294,144 +296,339 @@ class HomeDataDevice(RoborockBase): silent_ota_switch: bool | None = None setting: Any | None = None f: bool | None = None - device_features: DeviceFeatures | None = None - # seemingly not just str like I thought - example: '0000000000002000' and '0000000000002F63' - # def __post_init__(self): - # if self.feature_set is not None and self.new_feature_set is not None and self.new_feature_set != "": - # self.device_features = build_device_features(self.feature_set, self.new_feature_set) +class NewFeatureStrBit(IntEnum): + TWO_KEY_REAL_TIME_VIDEO = 32 + TWO_KEY_RTV_IN_CHARGING = 33 + DIRTY_REPLENISH_CLEAN = 34 + AUTO_DELIVERY_FIELD_IN_GLOBAL_STATUS = 35 + AVOID_COLLISION_MODE = 36 + VOICE_CONTROL = 37 + NEW_ENDPOINT = 38 + PUMPING_WATER = 39 + CORNER_MOP_STRECH = 40 + HOT_WASH_TOWEL = 41 + FLOOR_DIR_CLEAN_ANY_TIME = 42 + PET_SUPPLIES_DEEP_CLEAN = 43 + MOP_SHAKE_WATER_MAX = 45 + EXACT_CUSTOM_MODE = 47 + CARPET_CUSTOM_CLEAN = 49 + PET_SNAPSHOT = 50 + CUSTOM_CLEAN_MODE_COUNT = 51 + NEW_AI_RECOGNITION = 52 + AUTO_COLLECTION_2 = 53 + RIGHT_BRUSH_STRETCH = 54 + SMART_CLEAN_MODE_SET = 55 + DIRTY_OBJECT_DETECT = 56 + NO_NEED_CARPET_PRESS_SET = 57 + VOICE_CONTROL_LED = 58 + WATER_LEAK_CHECK = 60 + MIN_BATTERY_15_TO_CLEAN_TASK = 62 + GAP_DEEP_CLEAN = 63 + OBJECT_DETECT_CHECK = 64 + IDENTIFY_ROOM = 66 + MATTER = 67 + WORKDAY_HOLIDAY = 69 + CLEAN_DIRECT_STATUS = 70 + MAP_ERASER = 71 + OPTIMIZE_BATTERY = 72 + ACTIVATE_VIDEO_CHARGING_AND_STANDBY = 73 + CARPET_LONG_HAIRED = 75 + CLEAN_HISTORY_TIME_LINE = 76 + MAX_ZONE_OPENED = 77 + EXHIBITION_FUNCTION = 78 + LDS_LIFTING = 79 + AUTO_TEAR_DOWN_MOP = 80 + SAMLL_SIDE_MOP = 81 + SUPPORT_SIDE_BRUSH_UP_DOWN = 82 + DRY_INTERVAL_TIMER = 83 + UVC_STERILIZE = 84 + MIDWAY_BACK_TO_DOCK = 85 + SUPPORT_MAIN_BRUSH_UP_DOWN = 86 + EGG_DANCE_MODE = 87 @dataclass class DeviceFeatures(RoborockBase): - map_carpet_add_supported: bool - show_clean_finish_reason_supported: bool - resegment_supported: bool - video_monitor_supported: bool - any_state_transit_goto_supported: bool - fw_filter_obstacle_supported: bool - video_setting_supported: bool - ignore_unknown_map_object_supported: bool - set_child_supported: bool - carpet_supported: bool - mop_path_supported: bool - multi_map_segment_timer_supported: bool - custom_water_box_distance_supported: bool - wash_then_charge_cmd_supported: bool - room_name_supported: bool - current_map_restore_enabled: bool - photo_upload_supported: bool - shake_mop_set_supported: bool - map_beautify_internal_debug_supported: bool - new_data_for_clean_history: bool - new_data_for_clean_history_detail: bool - flow_led_setting_supported: bool - dust_collection_setting_supported: bool - rpc_retry_supported: bool - avoid_collision_supported: bool - support_set_switch_map_mode: bool - support_smart_scene: bool - support_floor_edit: bool - support_furniture: bool - support_room_tag: bool - support_quick_map_builder: bool - support_smart_global_clean_with_custom_mode: bool - record_allowed: bool - careful_slow_map_supported: bool - egg_mode_supported: bool - unsave_map_reason_supported: bool - carpet_show_on_map: bool - supported_valley_electricity: bool - drying_supported: bool - download_test_voice_supported: bool - support_backup_map: bool - support_custom_mode_in_cleaning: bool - support_remote_control_in_call: bool - support_set_volume_in_call: bool - support_clean_estimate: bool - support_custom_dnd: bool - carpet_deep_clean_supported: bool - stuck_zone_supported: bool - custom_door_sill_supported: bool - clean_route_fast_mode_supported: bool - cliff_zone_supported: bool - smart_door_sill_supported: bool - support_floor_direction: bool - wifi_manage_supported: bool - back_charge_auto_wash_supported: bool - support_incremental_map: bool - offline_map_supported: bool - - -def build_device_features(feature_set: str, new_feature_set: str) -> DeviceFeatures: - new_feature_set_int = int(new_feature_set) - feature_set_int = int(feature_set) - new_feature_set_divided = int(new_feature_set_int / (2**32)) - # Convert last 8 digits of new feature set into hexadecimal number - converted_new_feature_set = int("0x" + new_feature_set[-8:], 16) - new_feature_set_mod_8: bool = len(new_feature_set) % 8 == 0 - return DeviceFeatures( - map_carpet_add_supported=bool(1073741824 & new_feature_set_int), - show_clean_finish_reason_supported=bool(1 & new_feature_set_int), - resegment_supported=bool(4 & new_feature_set_int), - video_monitor_supported=bool(8 & new_feature_set_int), - any_state_transit_goto_supported=bool(16 & new_feature_set_int), - fw_filter_obstacle_supported=bool(32 & new_feature_set_int), - video_setting_supported=bool(64 & new_feature_set_int), - ignore_unknown_map_object_supported=bool(128 & new_feature_set_int), - set_child_supported=bool(256 & new_feature_set_int), - carpet_supported=bool(512 & new_feature_set_int), - mop_path_supported=bool(2048 & new_feature_set_int), - multi_map_segment_timer_supported=bool(feature_set_int and 4096 & new_feature_set_int), - custom_water_box_distance_supported=bool(new_feature_set_int and 2147483648 & new_feature_set_int), - wash_then_charge_cmd_supported=bool((new_feature_set_divided >> 5) & 1), - room_name_supported=bool(16384 & new_feature_set_int), - current_map_restore_enabled=bool(8192 & new_feature_set_int), - photo_upload_supported=bool(65536 & new_feature_set_int), - shake_mop_set_supported=bool(262144 & new_feature_set_int), - map_beautify_internal_debug_supported=bool(2097152 & new_feature_set_int), - new_data_for_clean_history=bool(4194304 & new_feature_set_int), - new_data_for_clean_history_detail=bool(8388608 & new_feature_set_int), - flow_led_setting_supported=bool(16777216 & new_feature_set_int), - dust_collection_setting_supported=bool(33554432 & new_feature_set_int), - rpc_retry_supported=bool(67108864 & new_feature_set_int), - avoid_collision_supported=bool(134217728 & new_feature_set_int), - support_set_switch_map_mode=bool(268435456 & new_feature_set_int), - support_smart_scene=bool(new_feature_set_divided & 2), - support_floor_edit=bool(new_feature_set_divided & 8), - support_furniture=bool((new_feature_set_divided >> 4) & 1), - support_room_tag=bool((new_feature_set_divided >> 6) & 1), - support_quick_map_builder=bool((new_feature_set_divided >> 7) & 1), - support_smart_global_clean_with_custom_mode=bool((new_feature_set_divided >> 8) & 1), - record_allowed=bool(1024 & new_feature_set_int), - careful_slow_map_supported=bool((new_feature_set_divided >> 9) & 1), - egg_mode_supported=bool((new_feature_set_divided >> 10) & 1), - unsave_map_reason_supported=bool((new_feature_set_divided >> 14) & 1), - carpet_show_on_map=bool((new_feature_set_divided >> 12) & 1), - supported_valley_electricity=bool((new_feature_set_divided >> 13) & 1), - # This one could actually be incorrect - # ((t.robotNewFeatures / 2 ** 32) >> 15) & 1 && (module422.DMM.isTopazSV_CE || 'cn' == t.deviceLocation)); - drying_supported=bool((new_feature_set_divided >> 15) & 1), - download_test_voice_supported=bool((new_feature_set_divided >> 16) & 1), - support_backup_map=bool((new_feature_set_divided >> 17) & 1), - support_custom_mode_in_cleaning=bool((new_feature_set_divided >> 18) & 1), - support_remote_control_in_call=bool((new_feature_set_divided >> 19) & 1), - support_set_volume_in_call=new_feature_set_mod_8 and bool(1 & converted_new_feature_set), - support_clean_estimate=new_feature_set_mod_8 and bool(2 & converted_new_feature_set), - support_custom_dnd=new_feature_set_mod_8 and bool(4 & converted_new_feature_set), - carpet_deep_clean_supported=bool(8 & converted_new_feature_set), - stuck_zone_supported=new_feature_set_mod_8 and bool(16 & converted_new_feature_set), - custom_door_sill_supported=new_feature_set_mod_8 and bool(32 & converted_new_feature_set), - clean_route_fast_mode_supported=bool(256 & converted_new_feature_set), - cliff_zone_supported=new_feature_set_mod_8 and bool(512 & converted_new_feature_set), - smart_door_sill_supported=new_feature_set_mod_8 and bool(1024 & converted_new_feature_set), - support_floor_direction=new_feature_set_mod_8 and bool(2048 & converted_new_feature_set), - wifi_manage_supported=bool(128 & converted_new_feature_set), - back_charge_auto_wash_supported=bool(4096 & converted_new_feature_set), - support_incremental_map=bool(8192 & converted_new_feature_set), - offline_map_supported=bool(16384 & converted_new_feature_set), - ) + """Represents the features supported by a Roborock device.""" + + # Features derived from robot_new_features + is_map_carpet_add_support: bool + is_show_clean_finish_reason_supported: bool + is_resegment_supported: bool + is_video_monitor_supported: bool + is_any_state_transit_goto_supported: bool + is_fw_filter_obstacle_supported: bool + is_video_settings_supported: bool + is_ignore_unknown_map_object_supported: bool + is_set_child_supported: bool + is_carpet_supported: bool + is_mop_path_supported: bool + is_multi_map_segment_timer_supported: bool + is_custom_water_box_distance_supported: bool + is_wash_then_charge_cmd_supported: bool + is_room_name_supported: bool + is_current_map_restore_enabled: bool + is_photo_upload_supported: bool + is_shake_mop_set_supported: bool + is_map_beautify_internal_debug_supported: bool + is_new_data_for_clean_history_supported: bool + is_new_data_for_clean_history_detail_supported: bool + is_flow_led_setting_supported: bool + is_dust_collection_setting_supported: bool + is_rpc_retry_supported: bool + is_avoid_collision_supported: bool + is_support_set_switch_map_mode_supported: bool + is_support_smart_scene_supported: bool + is_support_floor_edit_supported: bool + is_support_furniture_supported: bool + is_support_room_tag_supported: bool + is_support_quick_map_builder_supported: bool + is_support_smart_global_clean_with_custom_mode_supported: bool + is_record_allowed: bool + is_careful_slow_mop_supported: bool + is_egg_mode_supported: bool + is_carpet_show_on_map_supported: bool + is_supported_valley_electricity_supported: bool + is_unsave_map_reason_supported: bool + is_supported_drying_supported: bool + is_supported_download_test_voice_supported: bool + is_support_backup_map_supported: bool + is_support_custom_mode_in_cleaning_supported: bool + is_support_remote_control_in_call_supported: bool + + # Features derived from unhexed_feature_info + is_support_set_volume_in_call: bool + is_support_clean_estimate: bool + is_support_custom_dnd: bool + is_carpet_deep_clean_supported: bool + is_support_stuck_zone: bool + is_support_custom_door_sill: bool + is_wifi_manage_supported: bool + is_clean_route_fast_mode_supported: bool + is_support_cliff_zone: bool + is_support_smart_door_sill: bool + is_support_floor_direction: bool + is_back_charge_auto_wash_supported: bool + is_super_deep_wash_supported: bool + is_ces2022_supported: bool + is_dss_believable_supported: bool + is_main_brush_up_down_supported: bool + is_goto_pure_clean_path_supported: bool + is_water_up_down_drain_supported: bool + is_setting_carpet_first_supported: bool + is_clean_route_deep_slow_plus_supported: bool + is_left_water_drain_supported: bool + is_clean_count_setting_supported: bool + is_corner_clean_mode_supported: bool + + # --- Features from new_feature_info_str --- + is_two_key_real_time_video_supported: bool + is_two_key_rtv_in_charging_supported: bool + is_dirty_replenish_clean_supported: bool + is_avoid_collision_mode_str_supported: bool + is_voice_control_str_supported: bool + is_new_endpoint_supported: bool + is_corner_mop_strech_supported: bool + is_hot_wash_towel_supported: bool + is_floor_dir_clean_any_time_supported: bool + is_pet_supplies_deep_clean_supported: bool + is_mop_shake_water_max_supported: bool + is_custom_clean_mode_count_supported: bool + is_exact_custom_mode_supported: bool + is_carpet_custom_clean_supported: bool + is_pet_snapshot_supported: bool + is_new_ai_recognition_supported: bool + is_auto_collection_2_supported: bool + is_right_brush_stretch_supported: bool + is_smart_clean_mode_set_supported: bool + is_dirty_object_detect_supported: bool + is_no_need_carpet_press_set_supported: bool + is_voice_control_led_supported: bool + is_water_leak_check_supported: bool + is_min_battery_15_to_clean_task_supported: bool + is_gap_deep_clean_supported: bool + is_object_detect_check_supported: bool + is_identify_room_supported: bool + is_matter_supported: bool + + @classmethod + def _is_new_feature_str_support(cls, o: int, new_feature_info_str: str) -> bool: + """ + Checks feature 'o' in hex string 'new_feature_info_str'. + """ + try: + l = o % 4 + target_index = -((o // 4) + 1) + + p = new_feature_info_str[target_index] + + hex_char_value = int(p, 16) + + is_set = (hex_char_value >> l) & 1 + + return bool(is_set) + except Exception: + return False + + @classmethod + def from_feature_flags( + cls, robot_new_features: int, new_feature_set: str, product_nickname: RoborockProductNickname + ) -> DeviceFeatures: + """Creates a DeviceFeatures instance from raw feature flags.""" + unhexed_feature_info = int(new_feature_set[-8:], 16) + + upper_32_bits = robot_new_features // (2**32) + + return cls( + is_map_carpet_add_support=bool(1073741824 & robot_new_features), + is_show_clean_finish_reason_supported=bool(1 & robot_new_features), + is_resegment_supported=bool(4 & robot_new_features), + is_video_monitor_supported=bool(8 & robot_new_features), + is_any_state_transit_goto_supported=bool(16 & robot_new_features), + is_fw_filter_obstacle_supported=bool(32 & robot_new_features), + is_video_settings_supported=bool(64 & robot_new_features), + is_ignore_unknown_map_object_supported=bool(128 & robot_new_features), + is_set_child_supported=bool(256 & robot_new_features), + is_carpet_supported=bool(512 & robot_new_features), + is_mop_path_supported=bool(2048 & robot_new_features), + is_multi_map_segment_timer_supported=False, # TODO + is_custom_water_box_distance_supported=bool(2147483648 & robot_new_features), + is_wash_then_charge_cmd_supported=bool(robot_new_features and ((upper_32_bits >> 5) & 1)), + is_room_name_supported=bool(16384 & robot_new_features), + is_current_map_restore_enabled=bool(8192 & robot_new_features), + is_photo_upload_supported=bool(65536 & robot_new_features), + is_shake_mop_set_supported=bool(262144 & robot_new_features), + is_map_beautify_internal_debug_supported=bool(2097152 & robot_new_features), + is_new_data_for_clean_history_supported=bool(4194304 & robot_new_features), + is_new_data_for_clean_history_detail_supported=bool(8388608 & robot_new_features), + is_flow_led_setting_supported=bool(16777216 & robot_new_features), + is_dust_collection_setting_supported=bool(33554432 & robot_new_features), + is_rpc_retry_supported=bool(67108864 & robot_new_features), + is_avoid_collision_supported=bool(134217728 & robot_new_features), + is_support_set_switch_map_mode_supported=bool(268435456 & robot_new_features), + is_support_smart_scene_supported=bool(robot_new_features and (upper_32_bits & 2)), + is_support_floor_edit_supported=bool(robot_new_features and (upper_32_bits & 8)), + is_support_furniture_supported=bool(robot_new_features and ((upper_32_bits >> 4) & 1)), + is_support_room_tag_supported=bool(robot_new_features and ((upper_32_bits >> 6) & 1)), + is_support_quick_map_builder_supported=bool(robot_new_features and ((upper_32_bits >> 7) & 1)), + is_support_smart_global_clean_with_custom_mode_supported=bool( + robot_new_features and ((upper_32_bits >> 8) & 1) + ), + is_record_allowed=bool(1024 & robot_new_features), + is_careful_slow_mop_supported=bool(robot_new_features and ((upper_32_bits >> 9) & 1)), + is_egg_mode_supported=bool(robot_new_features and ((upper_32_bits >> 10) & 1)), + is_carpet_show_on_map_supported=bool(robot_new_features and ((upper_32_bits >> 12) & 1)), + is_supported_valley_electricity_supported=bool(robot_new_features and ((upper_32_bits >> 13) & 1)), + is_unsave_map_reason_supported=bool(robot_new_features and ((upper_32_bits >> 14) & 1)), + is_supported_drying_supported=False, # TODO + is_supported_download_test_voice_supported=bool(robot_new_features and ((upper_32_bits >> 16) & 1)), + is_support_backup_map_supported=bool(robot_new_features and ((upper_32_bits >> 17) & 1)), + is_support_custom_mode_in_cleaning_supported=bool(robot_new_features and ((upper_32_bits >> 18) & 1)), + is_support_remote_control_in_call_supported=bool(robot_new_features and ((upper_32_bits >> 19) & 1)), + # Features from unhexed_feature_info + is_support_set_volume_in_call=bool(1 & unhexed_feature_info), + is_support_clean_estimate=bool(2 & unhexed_feature_info), + is_support_custom_dnd=bool(4 & unhexed_feature_info), + is_carpet_deep_clean_supported=bool(8 & unhexed_feature_info), + is_support_stuck_zone=bool(16 & unhexed_feature_info), + is_support_custom_door_sill=bool(32 & unhexed_feature_info), + is_wifi_manage_supported=bool(128 & unhexed_feature_info), + is_clean_route_fast_mode_supported=bool(256 & unhexed_feature_info), + is_support_cliff_zone=bool(512 & unhexed_feature_info), + is_support_smart_door_sill=bool(1024 & unhexed_feature_info), + is_support_floor_direction=bool(2048 & unhexed_feature_info), + is_back_charge_auto_wash_supported=bool(4096 & unhexed_feature_info), + is_super_deep_wash_supported=bool(32768 & unhexed_feature_info), + is_ces2022_supported=bool(65536 & unhexed_feature_info), + is_dss_believable_supported=bool(131072 & unhexed_feature_info), + is_main_brush_up_down_supported=bool(262144 & unhexed_feature_info), + is_goto_pure_clean_path_supported=bool(524288 & unhexed_feature_info), + is_water_up_down_drain_supported=bool(1048576 & unhexed_feature_info), + is_setting_carpet_first_supported=bool(8388608 & unhexed_feature_info), + is_clean_route_deep_slow_plus_supported=bool(16777216 & unhexed_feature_info), + is_left_water_drain_supported=bool(134217728 & unhexed_feature_info), + is_clean_count_setting_supported=bool(1073741824 & unhexed_feature_info), + is_corner_clean_mode_supported=bool(2147483648 & unhexed_feature_info), + # Features from is_new_feature_str_support + is_two_key_real_time_video_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.TWO_KEY_REAL_TIME_VIDEO, new_feature_set + ), + is_two_key_rtv_in_charging_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.TWO_KEY_RTV_IN_CHARGING, new_feature_set + ), + is_dirty_replenish_clean_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.DIRTY_REPLENISH_CLEAN, new_feature_set + ), + is_avoid_collision_mode_str_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.AVOID_COLLISION_MODE, new_feature_set + ), + is_voice_control_str_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.VOICE_CONTROL, new_feature_set + ), + is_new_endpoint_supported=cls._is_new_feature_str_support(NewFeatureStrBit.NEW_ENDPOINT, new_feature_set), + is_corner_mop_strech_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.CORNER_MOP_STRECH, new_feature_set + ), + is_hot_wash_towel_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.HOT_WASH_TOWEL, new_feature_set + ), + is_floor_dir_clean_any_time_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.FLOOR_DIR_CLEAN_ANY_TIME, new_feature_set + ), + is_pet_supplies_deep_clean_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.PET_SUPPLIES_DEEP_CLEAN, new_feature_set + ), + is_mop_shake_water_max_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.MOP_SHAKE_WATER_MAX, new_feature_set + ), + is_custom_clean_mode_count_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.CUSTOM_CLEAN_MODE_COUNT, new_feature_set + ), + is_exact_custom_mode_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.EXACT_CUSTOM_MODE, new_feature_set + ), + is_carpet_custom_clean_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.CARPET_CUSTOM_CLEAN, new_feature_set + ), + is_pet_snapshot_supported=cls._is_new_feature_str_support(NewFeatureStrBit.PET_SNAPSHOT, new_feature_set), + is_new_ai_recognition_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.NEW_AI_RECOGNITION, new_feature_set + ), + is_auto_collection_2_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.AUTO_COLLECTION_2, new_feature_set + ), + is_right_brush_stretch_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.RIGHT_BRUSH_STRETCH, new_feature_set + ), + is_smart_clean_mode_set_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.SMART_CLEAN_MODE_SET, new_feature_set + ), + is_dirty_object_detect_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.DIRTY_OBJECT_DETECT, new_feature_set + ), + is_no_need_carpet_press_set_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.NO_NEED_CARPET_PRESS_SET, new_feature_set + ), + is_voice_control_led_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.VOICE_CONTROL_LED, new_feature_set + ), + is_water_leak_check_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.WATER_LEAK_CHECK, new_feature_set + ), + is_min_battery_15_to_clean_task_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.MIN_BATTERY_15_TO_CLEAN_TASK, new_feature_set + ), + is_gap_deep_clean_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.GAP_DEEP_CLEAN, new_feature_set + ), + is_object_detect_check_supported=cls._is_new_feature_str_support( + NewFeatureStrBit.OBJECT_DETECT_CHECK, new_feature_set + ), + is_identify_room_supported=cls._is_new_feature_str_support(NewFeatureStrBit.IDENTIFY_ROOM, new_feature_set), + is_matter_supported=cls._is_new_feature_str_support(NewFeatureStrBit.MATTER, new_feature_set), + ) @dataclass @@ -840,6 +1037,17 @@ class DeviceData(RoborockBase): device: HomeDataDevice model: str host: str | None = None + product_nickname: RoborockProductNickname | None = None + device_features: DeviceFeatures | None = None + + def __post_init__(self): + self.product_nickname = short_model_to_enum.get(self.model.split(".")[-1], RoborockProductNickname.PEARLPLUS) + robot_new_features = int(self.device.feature_set) if self.device.feature_set else 0 + self.device_features = DeviceFeatures.from_feature_flags( + robot_new_features, + self.device.new_feature_set if self.device.new_feature_set is not None else "00000000", + self.product_nickname, + ) @dataclass From f480b51c412ac502c029810c5277def41508d3ee Mon Sep 17 00:00:00 2001 From: Luke Date: Fri, 11 Apr 2025 10:49:21 -0400 Subject: [PATCH 4/8] chore: some misc changes --- roborock/device_trait.py | 79 ++++++++++++++++++++++++++++++ roborock/device_traits/__init__.py | 0 roborock/mqtt_manager.py | 3 +- roborock/roborock_device.py | 71 ++++++++++++++++++--------- roborock/roborock_message.py | 21 ++++---- 5 files changed, 140 insertions(+), 34 deletions(-) create mode 100644 roborock/device_trait.py create mode 100644 roborock/device_traits/__init__.py diff --git a/roborock/device_trait.py b/roborock/device_trait.py new file mode 100644 index 00000000..deb34884 --- /dev/null +++ b/roborock/device_trait.py @@ -0,0 +1,79 @@ +import datetime +from abc import ABC, abstractmethod +from collections.abc import Awaitable, Callable +from dataclasses import dataclass + +from . import RoborockCommand +from .containers import Consumable, DeviceFeatures, DnDTimer, RoborockBase + + +@dataclass +class DeviceTrait(ABC): + handle_command: RoborockCommand + _status_type: type[RoborockBase] = RoborockBase + + def __init__(self, send_command: Callable[..., Awaitable[None]]): + self.send_command = send_command + self.status: RoborockBase | None = None + self.subscriptions = [] + + @classmethod + @abstractmethod + def supported(cls, features: DeviceFeatures) -> bool: + raise NotImplementedError + + def on_message(self, data: dict) -> None: + self.status = self._status_type.from_dict(data) + for callback in self.subscriptions: + callback(self.status) + + def subscribe(self, callable: Callable): + # Maybe needs to handle async too? + self.subscriptions.append(callable) + + @abstractmethod + def get(self): + raise NotImplementedError + + +class DndTrait(DeviceTrait): + handle_command: RoborockCommand = RoborockCommand.GET_DND_TIMER + _status_type: type[DnDTimer] = DnDTimer + status: DnDTimer + + def __init__(self, send_command: Callable[..., Awaitable[None]]): + super().__init__(send_command) + + @classmethod + def supported(cls, features: DeviceFeatures) -> bool: + return features.is_support_custom_dnd + + async def update_dnd(self, enabled: bool, start_time: datetime.time, end_time: datetime.time) -> None: + if self.status.enabled and not enabled: + await self.send_command(RoborockCommand.CLOSE_DND_TIMER) + else: + start = start_time if start_time is not None else self.status.start_time + end = end_time if end_time is not None else self.status.end_time + await self.send_command(RoborockCommand.SET_DND_TIMER, [start.hour, start.minute, end.hour, end.minute]) + + async def get(self) -> None: + await self.send_command(RoborockCommand.GET_DND_TIMER) + + +class ConsumableTrait(DeviceTrait): + handle_command = RoborockCommand.GET_CONSUMABLE + _status_type: type[Consumable] = DnDTimer + status: Consumable + + def __init__(self, send_command: Callable[..., Awaitable[None]]): + super().__init__(send_command) + + @classmethod + def supported(cls, features: DeviceFeatures) -> bool: + return True + + async def reset_consumable(self, consumable: str) -> None: + await self.send_command(RoborockCommand.RESET_CONSUMABLE, [consumable]) + + async def get(self) -> None: + await self.send_command(RoborockCommand.GET_CONSUMABLE) diff --git a/roborock/device_traits/__init__.py b/roborock/device_traits/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/roborock/mqtt_manager.py b/roborock/mqtt_manager.py index c134c26f..1bcb62a9 100644 --- a/roborock/mqtt_manager.py +++ b/roborock/mqtt_manager.py @@ -89,7 +89,8 @@ async def unsubscribe(device: DeviceData): device_id = message.topic.value.split("/")[-1] device = device_map[device_id] message = MessageParser.parse(message.payload, device.device.local_key) - callbacks[device_id](message) + for m in message[0]: + callbacks[device_id](m) except Exception: ... diff --git a/roborock/roborock_device.py b/roborock/roborock_device.py index c54f9ed3..861dcde7 100644 --- a/roborock/roborock_device.py +++ b/roborock/roborock_device.py @@ -6,7 +6,8 @@ import time from . import RoborockCommand -from .containers import DeviceData, UserData +from .containers import DeviceData, ModelStatus, S7MaxVStatus, Status, UserData +from .device_trait import ConsumableTrait, DeviceTrait, DndTrait from .mqtt_manager import RoborockMqttManager from .protocol import MessageParser, Utils from .roborock_message import RoborockMessage, RoborockMessageProtocol @@ -25,26 +26,25 @@ def __init__(self, user_data: UserData, device_info: DeviceData): self._local_endpoint = "abc" self._nonce = secrets.token_bytes(16) self.manager = RoborockMqttManager() - self.update_commands = self.determine_supported_commands() - - def determine_supported_commands(self): - # All devices support these - supported_commands = { - RoborockCommand.GET_CONSUMABLE, - RoborockCommand.GET_STATUS, - RoborockCommand.GET_CLEAN_SUMMARY, - } - # Get what features we can from the feature_set info. - - # If a command is not described in feature_set, we should just add it anyways and then let it fail on the first call and remove it. - robot_new_features = int(self.device_info.device.feature_set) - new_feature_info_str = self.device_info.device.new_feature_set - if 33554432 & int(robot_new_features): - supported_commands.add(RoborockCommand.GET_DUST_COLLECTION_MODE) - if 2 & int(new_feature_info_str[-8:], 16): - # TODO: May not be needed as i think this can just be found in Status, but just POC - supported_commands.add(RoborockCommand.APP_GET_CLEAN_ESTIMATE_INFO) - return supported_commands + self._message_id_types: dict[int, DeviceTrait] = {} + self._command_to_trait = {} + self._all_supported_traits = [] + self._dnd_trait: DndTrait | None = self.determine_supported_traits(DndTrait) + self._consumable_trait: ConsumableTrait | None = self.determine_supported_traits(ConsumableTrait) + self._status_type: type[Status] = ModelStatus.get(device_info.model, S7MaxVStatus) + + def determine_supported_traits(self, trait: type[DeviceTrait]): + def _send_command( + method: RoborockCommand | str, params: list | dict | int | None = None, use_cloud: bool = True + ): + return self.send_message(method, params, use_cloud) + + if trait.supported(self.device_info.device_features): + trait_instance = trait(_send_command) + self._all_supported_traits.append(trait(_send_command)) + self._command_to_trait[trait.handle_command] = trait_instance + return trait_instance + return None async def connect(self): """Connect via MQTT and Local if possible.""" @@ -52,8 +52,8 @@ async def connect(self): await self.update() async def update(self): - for cmd in self.update_commands: - await self.send_message(method=cmd) + for trait in self._all_supported_traits: + await trait.get() def _get_payload( self, @@ -91,7 +91,9 @@ async def send_message( request_id, timestamp, payload = self._get_payload(method, params, True, use_cloud) request_protocol = RoborockMessageProtocol.RPC_REQUEST roborock_message = RoborockMessage(timestamp=timestamp, protocol=request_protocol, payload=payload) - + if request_id in self._message_id_types: + raise Exception("Duplicate id!") + self._message_id_types[request_id] = self._command_to_trait[method] local_key = self.device_info.device.local_key msg = MessageParser.build(roborock_message, local_key, False) if use_cloud: @@ -101,6 +103,19 @@ async def send_message( pass def on_message(self, message: RoborockMessage): + message_payload = message.get_payload() + message_id = message.get_request_id() + for data_point_number, data_point in message_payload.get("dps").items(): + if data_point_number == "102": + data_point_response = json.loads(data_point) + result = data_point_response.get("result") + if isinstance(result, list) and len(result) == 1: + result = result[0] + if result and (trait := self._message_id_types.get(message_id)) is not None: + trait.on_message(result) + if (error := result.get("error")) is not None: + print(error) + print() # If message is command not supported - remove from self.update_commands # If message is an error - log it? @@ -115,3 +130,11 @@ def on_message(self, message: RoborockMessage): # This should also probably be split with on_cloud_message and on_local_message. print(message) + + @property + def dnd(self) -> DndTrait | None: + return self._dnd_trait + + @property + def consumable(self) -> ConsumableTrait | None: + return self._consumable_trait diff --git a/roborock/roborock_message.py b/roborock/roborock_message.py index 1774c3bb..943e953d 100644 --- a/roborock/roborock_message.py +++ b/roborock/roborock_message.py @@ -161,11 +161,16 @@ class RoborockMessage: random: int = field(default_factory=lambda: get_next_int(10000, 99999)) timestamp: int = field(default_factory=lambda: math.floor(time.time())) message_retry: MessageRetry | None = None + _parsed_payload: dict | None = None + + def get_payload(self) -> dict | None: + if self.payload and not self._parsed_payload: + self._parsed_payload = json.loads(self.payload.decode()) + return self._parsed_payload def get_request_id(self) -> int | None: - if self.payload: - payload = json.loads(self.payload.decode()) - for data_point_number, data_point in payload.get("dps").items(): + if self._parsed_payload: + for data_point_number, data_point in self._parsed_payload.get("dps").items(): if data_point_number in ["101", "102"]: data_point_response = json.loads(data_point) return data_point_response.get("id") @@ -180,9 +185,8 @@ def get_method(self) -> str | None: if self.message_retry: return self.message_retry.method protocol = self.protocol - if self.payload and protocol in [4, 5, 101, 102]: - payload = json.loads(self.payload.decode()) - for data_point_number, data_point in payload.get("dps").items(): + if self._parsed_payload and protocol in [4, 5, 101, 102]: + for data_point_number, data_point in self._parsed_payload.get("dps").items(): if data_point_number in ["101", "102"]: data_point_response = json.loads(data_point) return data_point_response.get("method") @@ -190,9 +194,8 @@ def get_method(self) -> str | None: def get_params(self) -> list | dict | None: protocol = self.protocol - if self.payload and protocol in [4, 101, 102]: - payload = json.loads(self.payload.decode()) - for data_point_number, data_point in payload.get("dps").items(): + if self._parsed_payload and protocol in [4, 101, 102]: + for data_point_number, data_point in self._parsed_payload.get("dps").items(): if data_point_number in ["101", "102"]: data_point_response = json.loads(data_point) return data_point_response.get("params") From 51ad2441dd6aa548271a48919cd2beaddb095c7d Mon Sep 17 00:00:00 2001 From: Luke Date: Tue, 6 May 2025 08:29:07 -0400 Subject: [PATCH 5/8] chore: random changes --- roborock/mqtt_manager.py | 114 ----------------------------------- roborock/roborock_device.py | 28 +++++++-- roborock/roborock_message.py | 9 ++- 3 files changed, 30 insertions(+), 121 deletions(-) delete mode 100644 roborock/mqtt_manager.py diff --git a/roborock/mqtt_manager.py b/roborock/mqtt_manager.py deleted file mode 100644 index 1bcb62a9..00000000 --- a/roborock/mqtt_manager.py +++ /dev/null @@ -1,114 +0,0 @@ -from __future__ import annotations - -import asyncio -import dataclasses -import logging -from collections.abc import Coroutine -from typing import Callable, Self -from urllib.parse import urlparse - -import aiomqtt -from aiomqtt import TLSParameters - -from roborock import RoborockException, UserData -from roborock.protocol import MessageParser, md5hex - -from .containers import DeviceData - -LOGGER = logging.getLogger(__name__) - - -@dataclasses.dataclass -class ClientWrapper: - publish_function: Coroutine[None] - unsubscribe_function: Coroutine[None] - subscribe_function: Coroutine[None] - - -class RoborockMqttManager: - client_wrappers: dict[str, ClientWrapper] = {} - _instance: Self = None - - def __new__(cls) -> RoborockMqttManager: - if cls._instance is None: - cls._instance = super().__new__(cls) - return cls._instance - - async def connect(self, user_data: UserData): - # Add some kind of lock so we don't try to connect if we are already trying to connect the same account. - if user_data.rriot.u not in self.client_wrappers: - loop = asyncio.get_event_loop() - loop.create_task(self._new_connect(user_data)) - - async def _new_connect(self, user_data: UserData): - rriot = user_data.rriot - mqtt_user = rriot.u - hashed_user = md5hex(mqtt_user + ":" + rriot.k)[2:10] - url = urlparse(rriot.r.m) - if not isinstance(url.hostname, str): - raise RoborockException("Url parsing returned an invalid hostname") - mqtt_host = str(url.hostname) - mqtt_port = url.port - - mqtt_password = rriot.s - hashed_password = md5hex(mqtt_password + ":" + rriot.k)[16:] - LOGGER.debug("Connecting to %s for %s", mqtt_host, mqtt_user) - - async with aiomqtt.Client( - hostname=mqtt_host, - port=mqtt_port, - username=hashed_user, - password=hashed_password, - keepalive=60, - tls_params=TLSParameters(), - ) as client: - # TODO: Handle logic for when client loses connection - LOGGER.info("Connected to %s for %s", mqtt_host, mqtt_user) - callbacks: dict[str, Callable] = {} - device_map = {} - - async def publish(device: DeviceData, payload: bytes): - await client.publish(f"rr/m/i/{mqtt_user}/{hashed_user}/{device.device.duid}", payload=payload) - - async def subscribe(device: DeviceData, callback): - LOGGER.debug(f"Subscribing to rr/m/o/{mqtt_user}/{hashed_user}/{device.device.duid}") - await client.subscribe(f"rr/m/o/{mqtt_user}/{hashed_user}/{device.device.duid}") - LOGGER.debug(f"Subscribed to rr/m/o/{mqtt_user}/{hashed_user}/{device.device.duid}") - callbacks[device.device.duid] = callback - device_map[device.device.duid] = device - return - - async def unsubscribe(device: DeviceData): - await client.unsubscribe(f"rr/m/o/{mqtt_user}/{hashed_user}/{device.device.duid}") - - self.client_wrappers[user_data.rriot.u] = ClientWrapper( - publish_function=publish, unsubscribe_function=unsubscribe, subscribe_function=subscribe - ) - async for message in client.messages: - try: - device_id = message.topic.value.split("/")[-1] - device = device_map[device_id] - message = MessageParser.parse(message.payload, device.device.local_key) - for m in message[0]: - callbacks[device_id](m) - except Exception: - ... - - async def disconnect(self, user_data: UserData): - await self.client_wrappers[user_data.rriot.u].disconnect() - - async def subscribe(self, user_data: UserData, device: DeviceData, callback): - if user_data.rriot.u not in self.client_wrappers: - await self.connect(user_data) - # add some kind of lock to make sure we don't subscribe until the connection is successful - await asyncio.sleep(2) - await self.client_wrappers[user_data.rriot.u].subscribe_function(device, callback) - - async def unsubscribe(self): - pass - - async def publish(self, user_data: UserData, device, payload: bytes): - LOGGER.debug("Publishing topic for %s, Message: %s", device.device.duid, payload) - if user_data.rriot.u not in self.client_wrappers: - await self.connect(user_data) - await self.client_wrappers[user_data.rriot.u].publish_function(device, payload) diff --git a/roborock/roborock_device.py b/roborock/roborock_device.py index 861dcde7..2ea36385 100644 --- a/roborock/roborock_device.py +++ b/roborock/roborock_device.py @@ -4,12 +4,13 @@ import math import secrets import time +from urllib.parse import urlparse from . import RoborockCommand from .containers import DeviceData, ModelStatus, S7MaxVStatus, Status, UserData from .device_trait import ConsumableTrait, DeviceTrait, DndTrait -from .mqtt_manager import RoborockMqttManager -from .protocol import MessageParser, Utils +from .mqtt.roborock_session import MqttParams, RoborockMqttSession +from .protocol import MessageParser, Utils, md5hex from .roborock_message import RoborockMessage, RoborockMessageProtocol from .util import RoborockLoggerAdapter, get_next_int @@ -17,15 +18,25 @@ class RoborockDevice: + _mqtt_sessions: dict[str, RoborockMqttSession] = {} + def __init__(self, user_data: UserData, device_info: DeviceData): self.user_data = user_data self.device_info = device_info self.data = None self._logger = RoborockLoggerAdapter(device_info.device.name, _LOGGER) self._mqtt_endpoint = base64.b64encode(Utils.md5(user_data.rriot.k.encode())[8:14]).decode() + rriot = user_data.rriot + self._mqtt_user = rriot.u + self._hashed_user = md5hex(self._mqtt_user + ":" + rriot.k)[2:10] + url = urlparse(rriot.r.m) + self._mqtt_host = str(url) + self._mqtt_port = url.port + mqtt_password = rriot.s + self._hashed_password = md5hex(mqtt_password + ":" + rriot.k)[16:] + self._local_endpoint = "abc" self._nonce = secrets.token_bytes(16) - self.manager = RoborockMqttManager() self._message_id_types: dict[int, DeviceTrait] = {} self._command_to_trait = {} self._all_supported_traits = [] @@ -48,6 +59,14 @@ def _send_command( async def connect(self): """Connect via MQTT and Local if possible.""" + + MqttParams( + host=self._mqtt_host, + port=self._mqtt_port, + tls=True, + username=self._hashed_user, + password=self._hashed_password, + ) await self.manager.subscribe(self.user_data, self.device_info, self.on_message) await self.update() @@ -93,7 +112,8 @@ async def send_message( roborock_message = RoborockMessage(timestamp=timestamp, protocol=request_protocol, payload=payload) if request_id in self._message_id_types: raise Exception("Duplicate id!") - self._message_id_types[request_id] = self._command_to_trait[method] + if method in self._command_to_trait: + self._message_id_types[request_id] = self._command_to_trait[method] local_key = self.device_info.device.local_key msg = MessageParser.build(roborock_message, local_key, False) if use_cloud: diff --git a/roborock/roborock_message.py b/roborock/roborock_message.py index 943e953d..1e8535b0 100644 --- a/roborock/roborock_message.py +++ b/roborock/roborock_message.py @@ -169,7 +169,8 @@ def get_payload(self) -> dict | None: return self._parsed_payload def get_request_id(self) -> int | None: - if self._parsed_payload: + payload = self.get_payload() + if payload: for data_point_number, data_point in self._parsed_payload.get("dps").items(): if data_point_number in ["101", "102"]: data_point_response = json.loads(data_point) @@ -185,7 +186,8 @@ def get_method(self) -> str | None: if self.message_retry: return self.message_retry.method protocol = self.protocol - if self._parsed_payload and protocol in [4, 5, 101, 102]: + payload = self.get_payload() + if payload and protocol in [4, 5, 101, 102]: for data_point_number, data_point in self._parsed_payload.get("dps").items(): if data_point_number in ["101", "102"]: data_point_response = json.loads(data_point) @@ -194,7 +196,8 @@ def get_method(self) -> str | None: def get_params(self) -> list | dict | None: protocol = self.protocol - if self._parsed_payload and protocol in [4, 101, 102]: + payload = self.get_payload() + if payload and protocol in [4, 101, 102]: for data_point_number, data_point in self._parsed_payload.get("dps").items(): if data_point_number in ["101", "102"]: data_point_response = json.loads(data_point) From b214d5b59d46ccfc6a670e27f03ecbf76378ca1a Mon Sep 17 00:00:00 2001 From: Luke Date: Sat, 10 May 2025 16:10:54 -0400 Subject: [PATCH 6/8] chore: make things functional again --- roborock/containers.py | 5 ++ roborock/roborock_device.py | 91 +++++++++++++++++++------------------ 2 files changed, 52 insertions(+), 44 deletions(-) diff --git a/roborock/containers.py b/roborock/containers.py index d0be532f..5502aaa8 100644 --- a/roborock/containers.py +++ b/roborock/containers.py @@ -1049,6 +1049,11 @@ def __post_init__(self): self.product_nickname, ) + @property + def duid(self) -> str: + """Get the duid of the device.""" + return self.device.duid + @dataclass class RoomMapping(RoborockBase): diff --git a/roborock/roborock_device.py b/roborock/roborock_device.py index 2ea36385..7fff112b 100644 --- a/roborock/roborock_device.py +++ b/roborock/roborock_device.py @@ -27,13 +27,9 @@ def __init__(self, user_data: UserData, device_info: DeviceData): self._logger = RoborockLoggerAdapter(device_info.device.name, _LOGGER) self._mqtt_endpoint = base64.b64encode(Utils.md5(user_data.rriot.k.encode())[8:14]).decode() rriot = user_data.rriot - self._mqtt_user = rriot.u - self._hashed_user = md5hex(self._mqtt_user + ":" + rriot.k)[2:10] + hashed_user = md5hex(rriot.u + ":" + rriot.k)[2:10] url = urlparse(rriot.r.m) - self._mqtt_host = str(url) - self._mqtt_port = url.port mqtt_password = rriot.s - self._hashed_password = md5hex(mqtt_password + ":" + rriot.k)[16:] self._local_endpoint = "abc" self._nonce = secrets.token_bytes(16) @@ -43,6 +39,18 @@ def __init__(self, user_data: UserData, device_info: DeviceData): self._dnd_trait: DndTrait | None = self.determine_supported_traits(DndTrait) self._consumable_trait: ConsumableTrait | None = self.determine_supported_traits(ConsumableTrait) self._status_type: type[Status] = ModelStatus.get(device_info.model, S7MaxVStatus) + # TODO: One per client EVER + self.session = RoborockMqttSession( + MqttParams( + host=str(url.hostname), + port=url.port, + tls=True, + username=hashed_user, + password=md5hex(rriot.s + ":" + rriot.k)[16:], + ) + ) + self.input_topic = f"rr/m/i/{rriot.u}/{hashed_user}/{device_info.duid}" + self.output_topic = f"rr/m/o/{rriot.u}/{hashed_user}/{device_info.duid}" def determine_supported_traits(self, trait: type[DeviceTrait]): def _send_command( @@ -59,16 +67,9 @@ def _send_command( async def connect(self): """Connect via MQTT and Local if possible.""" - - MqttParams( - host=self._mqtt_host, - port=self._mqtt_port, - tls=True, - username=self._hashed_user, - password=self._hashed_password, - ) - await self.manager.subscribe(self.user_data, self.device_info, self.on_message) - await self.update() + if not self.session.connected: + await self.session.start() + await self.session.subscribe(self.output_topic, callback=self.on_message) async def update(self): for trait in self._all_supported_traits: @@ -117,39 +118,41 @@ async def send_message( local_key = self.device_info.device.local_key msg = MessageParser.build(roborock_message, local_key, False) if use_cloud: - await self.manager.publish(self.user_data, self.device_info, msg) + await self.session.publish(self.input_topic, msg) else: # Handle doing local commands pass - def on_message(self, message: RoborockMessage): - message_payload = message.get_payload() - message_id = message.get_request_id() - for data_point_number, data_point in message_payload.get("dps").items(): - if data_point_number == "102": - data_point_response = json.loads(data_point) - result = data_point_response.get("result") - if isinstance(result, list) and len(result) == 1: - result = result[0] - if result and (trait := self._message_id_types.get(message_id)) is not None: - trait.on_message(result) - if (error := result.get("error")) is not None: - print(error) - print() - # If message is command not supported - remove from self.update_commands - - # If message is an error - log it? - - # If message is 'ok' - ignore it - - # If message is anything else - store ids, and map back to id to determine message type. - # Then update self.data - - # If we haven't received a message in X seconds, the device is likely offline. I think we can continue the connection, - # but we should have some way to mark ourselves as unavailable. - - # This should also probably be split with on_cloud_message and on_local_message. - print(message) + def on_message(self, message_bytes: bytes): + messages = MessageParser.parse(message_bytes, self.device_info.device.local_key)[0] + for message in messages: + message_payload = message.get_payload() + message_id = message.get_request_id() + for data_point_number, data_point in message_payload.get("dps").items(): + if data_point_number == "102": + data_point_response = json.loads(data_point) + result = data_point_response.get("result") + if isinstance(result, list) and len(result) == 1: + result = result[0] + if result and (trait := self._message_id_types.get(message_id)) is not None: + trait.on_message(result) + if (error := result.get("error")) is not None: + print(error) + print() + # If message is command not supported - remove from self.update_commands + + # If message is an error - log it? + + # If message is 'ok' - ignore it + + # If message is anything else - store ids, and map back to id to determine message type. + # Then update self.data + + # If we haven't received a message in X seconds, the device is likely offline. I think we can continue the connection, + # but we should have some way to mark ourselves as unavailable. + + # This should also probably be split with on_cloud_message and on_local_message. + print(message) @property def dnd(self) -> DndTrait | None: From fbf859b2fc6787f1e9af42ac8a29e3f6d68c7c33 Mon Sep 17 00:00:00 2001 From: Luke Date: Fri, 30 May 2025 11:48:14 -0400 Subject: [PATCH 7/8] chore: change a bit of how traits work --- roborock/containers.py | 9 ++++ roborock/device_trait.py | 68 ++++++++++-------------------- roborock/device_traits/__init__.py | 1 + roborock/device_traits/dnd.py | 56 ++++++++++++++++++++++++ roborock/roborock_device.py | 15 ++----- 5 files changed, 93 insertions(+), 56 deletions(-) create mode 100644 roborock/device_traits/dnd.py diff --git a/roborock/containers.py b/roborock/containers.py index 5502aaa8..9acdb0f0 100644 --- a/roborock/containers.py +++ b/roborock/containers.py @@ -1184,3 +1184,12 @@ class DyadSndState(RoborockBase): @dataclass class DyadOtaNfo(RoborockBase): mqttOtaData: dict + + +@dataclass +class DndActions(RoborockBase): + dry: int | None = None + dust: int | None = None + led: int | None = None + resume: int | None = None + vol: int | None = None diff --git a/roborock/device_trait.py b/roborock/device_trait.py index deb34884..f8837fd0 100644 --- a/roborock/device_trait.py +++ b/roborock/device_trait.py @@ -1,16 +1,14 @@ -import datetime from abc import ABC, abstractmethod from collections.abc import Awaitable, Callable from dataclasses import dataclass from . import RoborockCommand -from .containers import Consumable, DeviceFeatures, DnDTimer, RoborockBase +from .containers import DeviceFeatures, RoborockBase @dataclass class DeviceTrait(ABC): handle_command: RoborockCommand - _status_type: type[RoborockBase] = RoborockBase def __init__(self, send_command: Callable[..., Awaitable[None]]): self.send_command = send_command @@ -22,8 +20,12 @@ def __init__(self, send_command: Callable[..., Awaitable[None]]): def supported(cls, features: DeviceFeatures) -> bool: raise NotImplementedError + @abstractmethod + def from_dict(cls, data: dict) -> bool: + raise NotImplementedError + def on_message(self, data: dict) -> None: - self.status = self._status_type.from_dict(data) + self.status = self.from_dict(data) for callback in self.subscriptions: callback(self.status) @@ -36,44 +38,20 @@ def get(self): raise NotImplementedError -class DndTrait(DeviceTrait): - handle_command: RoborockCommand = RoborockCommand.GET_DND_TIMER - _status_type: type[DnDTimer] = DnDTimer - status: DnDTimer - - def __init__(self, send_command: Callable[..., Awaitable[None]]): - super().__init__(send_command) - - @classmethod - def supported(cls, features: DeviceFeatures) -> bool: - return features.is_support_custom_dnd - - async def update_dnd(self, enabled: bool, start_time: datetime.time, end_time: datetime.time) -> None: - if self.status.enabled and not enabled: - await self.send_command(RoborockCommand.CLOSE_DND_TIMER) - else: - start = start_time if start_time is not None else self.status.start_time - end = end_time if end_time is not None else self.status.end_time - await self.send_command(RoborockCommand.SET_DND_TIMER, [start.hour, start.minute, end.hour, end.minute]) - - async def get(self) -> None: - await self.send_command(RoborockCommand.GET_DND_TIMER) - - -class ConsumableTrait(DeviceTrait): - handle_command = RoborockCommand.GET_CONSUMABLE - _status_type: type[Consumable] = DnDTimer - status: Consumable - - def __init__(self, send_command: Callable[..., Awaitable[None]]): - super().__init__(send_command) - - @classmethod - def supported(cls, features: DeviceFeatures) -> bool: - return True - - async def reset_consumable(self, consumable: str) -> None: - await self.send_command(RoborockCommand.RESET_CONSUMABLE, [consumable]) - - async def get(self) -> None: - await self.send_command(RoborockCommand.GET_CONSUMABLE) +# class ConsumableTrait(DeviceTrait): +# handle_command = RoborockCommand.GET_CONSUMABLE +# _status_type: type[Consumable] = DnDTimer +# status: Consumable +# +# def __init__(self, send_command: Callable[..., Awaitable[None]]): +# super().__init__(send_command) +# +# @classmethod +# def supported(cls, features: DeviceFeatures) -> bool: +# return True +# +# async def reset_consumable(self, consumable: str) -> None: +# await self.send_command(RoborockCommand.RESET_CONSUMABLE, [consumable]) +# +# async def get(self) -> None: +# await self.send_command(RoborockCommand.GET_CONSUMABLE) diff --git a/roborock/device_traits/__init__.py b/roborock/device_traits/__init__.py index e69de29b..7adf22ed 100644 --- a/roborock/device_traits/__init__.py +++ b/roborock/device_traits/__init__.py @@ -0,0 +1 @@ +from .dnd import Dnd diff --git a/roborock/device_traits/dnd.py b/roborock/device_traits/dnd.py new file mode 100644 index 00000000..08317788 --- /dev/null +++ b/roborock/device_traits/dnd.py @@ -0,0 +1,56 @@ +import datetime +from collections.abc import Awaitable, Callable + +from roborock import DeviceFeatures, DndActions, RoborockCommand +from roborock.device_trait import DeviceTrait + + +class Dnd(DeviceTrait): + handle_command: RoborockCommand = RoborockCommand.GET_DND_TIMER + + def __init__(self, send_command: Callable[..., Awaitable[None]]): + self.start_hour: int | None = None + self.start_minute: int | None = None + self.end_hour: int | None = None + self.end_minute: int | None = None + self.enabled: bool | None = None + self.start_time: datetime.time | None = None + self.end_time: datetime.time | None = None + self.actions: DndActions | None = None + super().__init__(send_command) + + def from_dict(self, dnd_dict: dict): + self.start_hour = dnd_dict.get("start_hour") + self.start_minute = dnd_dict.get("start_minute") + self.end_hour = dnd_dict.get("end_hour") + self.end_minute = dnd_dict.get("end_minute") + self.enabled = bool(dnd_dict.get("enabled")) + self.actions = DndActions.from_dict(dnd_dict.get("actions")) + self.start_time = ( + datetime.time(hour=self.start_hour, minute=self.start_minute) + if self.start_hour is not None and self.start_minute is not None + else None + ) + self.end_time = ( + datetime.time(hour=self.end_hour, minute=self.end_minute) + if self.end_hour is not None and self.end_minute is not None + else None + ) + + def to_dict(self) -> dict: + return {} + + @classmethod + def supported(cls, features: DeviceFeatures) -> bool: + return features.is_support_custom_dnd + + async def update_dnd(self, enabled: bool, start_time: datetime.time, end_time: datetime.time) -> None: + if self.enabled and not enabled: + await self.send_command(RoborockCommand.CLOSE_DND_TIMER) + else: + start = start_time if start_time is not None else self.start_time + end = end_time if end_time is not None else self.end_time + await self.send_command(RoborockCommand.SET_DND_TIMER, [start.hour, start.minute, end.hour, end.minute]) + + async def get(self) -> None: + await self.send_command(RoborockCommand.GET_DND_TIMER) diff --git a/roborock/roborock_device.py b/roborock/roborock_device.py index 7fff112b..e29b5916 100644 --- a/roborock/roborock_device.py +++ b/roborock/roborock_device.py @@ -8,7 +8,8 @@ from . import RoborockCommand from .containers import DeviceData, ModelStatus, S7MaxVStatus, Status, UserData -from .device_trait import ConsumableTrait, DeviceTrait, DndTrait +from .device_trait import DeviceTrait +from .device_traits import Dnd from .mqtt.roborock_session import MqttParams, RoborockMqttSession from .protocol import MessageParser, Utils, md5hex from .roborock_message import RoborockMessage, RoborockMessageProtocol @@ -36,8 +37,8 @@ def __init__(self, user_data: UserData, device_info: DeviceData): self._message_id_types: dict[int, DeviceTrait] = {} self._command_to_trait = {} self._all_supported_traits = [] - self._dnd_trait: DndTrait | None = self.determine_supported_traits(DndTrait) - self._consumable_trait: ConsumableTrait | None = self.determine_supported_traits(ConsumableTrait) + self._dnd_trait: Dnd | None = self.determine_supported_traits(Dnd) + # self._consumable_trait: ConsumableTrait | None = self.determine_supported_traits(ConsumableTrait) self._status_type: type[Status] = ModelStatus.get(device_info.model, S7MaxVStatus) # TODO: One per client EVER self.session = RoborockMqttSession( @@ -153,11 +154,3 @@ def on_message(self, message_bytes: bytes): # This should also probably be split with on_cloud_message and on_local_message. print(message) - - @property - def dnd(self) -> DndTrait | None: - return self._dnd_trait - - @property - def consumable(self) -> ConsumableTrait | None: - return self._consumable_trait From d98b72e462f956db7bf00aaec38506469710825f Mon Sep 17 00:00:00 2001 From: Luke Date: Sun, 15 Jun 2025 16:56:20 -0400 Subject: [PATCH 8/8] chore: some PR updates --- roborock/code_mappings.py | 193 ++++++----------- roborock/containers.py | 399 +++++++++++++---------------------- roborock/device_trait.py | 7 +- roborock/roborock_device.py | 1 - roborock/roborock_message.py | 21 +- 5 files changed, 220 insertions(+), 401 deletions(-) diff --git a/roborock/code_mappings.py b/roborock/code_mappings.py index 140bf0d8..8d1123b5 100644 --- a/roborock/code_mappings.py +++ b/roborock/code_mappings.py @@ -1,7 +1,8 @@ from __future__ import annotations import logging -from enum import Enum, IntEnum, StrEnum +from collections import namedtuple +from enum import Enum, IntEnum _LOGGER = logging.getLogger(__name__) completed_warnings = set() @@ -50,142 +51,66 @@ def items(cls: type[RoborockEnum]): return cls.as_dict().items() -class RoborockProductNickname(StrEnum): - """Enumeration of product nicknames.""" - - CORAL = "Coral" - CORALPRO = "CoralPro" - PEARL = "Pearl" - PEARLC = "PearlC" - PEARLE = "PearlE" - PEARLELITE = "PearlELite" - PEARLPLUS = "PearlPlus" - PEARLPLUSS = "PearlPlusS" - PEARLS = "PearlS" - PEARLSLITE = "PearlSLite" - RUBYPLUS = "RubyPlus" - RUBYSC = "RubySC" - RUBYSE = "RubySE" - RUBYSLITE = "RubySLite" - TANOS = "Tanos" - TANOSE = "TanosE" - TANOSS = "TanosS" - TANOSSC = "TanosSC" - TANOSSE = "TanosSE" - TANOSSMAX = "TanosSMax" - TANOSSLITE = "TanosSLite" - TANOSSPLUS = "TanosSPlus" - TANOSV = "TanosV" - TOPAZS = "TopazS" - TOPAZSC = "TopazSC" - TOPAZSPLUS = "TopazSPlus" - TOPAZSPOWER = "TopazSPower" - TOPAZSV = "TopazSV" - ULTRON = "Ultron" - ULTRONE = "UltronE" - ULTRONLITE = "UltronLite" - ULTRONSC = "UltronSC" - ULTRONSE = "UltronSE" - ULTRONSPLUS = "UltronSPlus" - ULTRONSV = "UltronSV" - VERDELITE = "Verdelite" - VIVIAN = "Vivian" - VIVIANC = "VivianC" - - -short_model_to_enum = { - # Pearl Series - "a103": RoborockProductNickname.PEARLC, - "a104": RoborockProductNickname.PEARLC, - "a116": RoborockProductNickname.PEARLPLUSS, - "a117": RoborockProductNickname.PEARLPLUSS, - "a136": RoborockProductNickname.PEARLPLUSS, - "a122": RoborockProductNickname.PEARLSLITE, - "a123": RoborockProductNickname.PEARLSLITE, - "a167": RoborockProductNickname.PEARLE, - "a168": RoborockProductNickname.PEARLE, - "a169": RoborockProductNickname.PEARLELITE, - "a170": RoborockProductNickname.PEARLELITE, - "a74": RoborockProductNickname.PEARL, - "a75": RoborockProductNickname.PEARL, - "a100": RoborockProductNickname.PEARLS, - "a101": RoborockProductNickname.PEARLS, - "a86": RoborockProductNickname.PEARLPLUS, - "a87": RoborockProductNickname.PEARLPLUS, - # Vivian Series - "a158": RoborockProductNickname.VIVIANC, - "a159": RoborockProductNickname.VIVIANC, - "a134": RoborockProductNickname.VIVIAN, - "a135": RoborockProductNickname.VIVIAN, - "a155": RoborockProductNickname.VIVIAN, - "a156": RoborockProductNickname.VIVIAN, +ProductInfo = namedtuple("ProductInfo", ["nickname", "short_models"]) + + +class RoborockProductNickname(Enum): # Coral Series - "a143": RoborockProductNickname.CORALPRO, - "a144": RoborockProductNickname.CORALPRO, - "a20": RoborockProductNickname.CORAL, - "a21": RoborockProductNickname.CORAL, + CORAL = ProductInfo(nickname="Coral", short_models=("a20", "a21")) + CORALPRO = ProductInfo(nickname="CoralPro", short_models=("a143", "a144")) + + # Pearl Series + PEARL = ProductInfo(nickname="Pearl", short_models=("a74", "a75")) + PEARLC = ProductInfo(nickname="PearlC", short_models=("a103", "a104")) + PEARLE = ProductInfo(nickname="PearlE", short_models=("a167", "a168")) + PEARLELITE = ProductInfo(nickname="PearlELite", short_models=("a169", "a170")) + PEARLPLUS = ProductInfo(nickname="PearlPlus", short_models=("a86", "a87")) + PEARLPLUSS = ProductInfo(nickname="PearlPlusS", short_models=("a116", "a117", "a136")) + PEARLS = ProductInfo(nickname="PearlS", short_models=("a100", "a101")) + PEARLSLITE = ProductInfo(nickname="PearlSLite", short_models=("a122", "a123")) + + # Ruby Series + RUBYPLUS = ProductInfo(nickname="RubyPlus", short_models=("t4", "s4")) + RUBYSC = ProductInfo(nickname="RubySC", short_models=("p5", "a08")) + RUBYSE = ProductInfo(nickname="RubySE", short_models=("a19",)) + RUBYSLITE = ProductInfo(nickname="RubySLite", short_models=("p6", "s5e", "a05")) + + # Tanos Series + TANOS = ProductInfo(nickname="Tanos", short_models=("t6", "s6")) + TANOSE = ProductInfo(nickname="TanosE", short_models=("t7", "a11")) + TANOSS = ProductInfo(nickname="TanosS", short_models=("a14", "a15")) + TANOSSC = ProductInfo(nickname="TanosSC", short_models=("a39", "a40")) + TANOSSE = ProductInfo(nickname="TanosSE", short_models=("a33", "a34")) + TANOSSMAX = ProductInfo(nickname="TanosSMax", short_models=("a52",)) + TANOSSLITE = ProductInfo(nickname="TanosSLite", short_models=("a37", "a38")) + TANOSSPLUS = ProductInfo(nickname="TanosSPlus", short_models=("a23", "a24")) + TANOSV = ProductInfo(nickname="TanosV", short_models=("t7p", "a09", "a10")) + + # Topaz Series + TOPAZS = ProductInfo(nickname="TopazS", short_models=("a29", "a30", "a76")) + TOPAZSC = ProductInfo(nickname="TopazSC", short_models=("a64", "a65")) + TOPAZSPLUS = ProductInfo(nickname="TopazSPlus", short_models=("a46", "a47", "a66")) + TOPAZSPOWER = ProductInfo(nickname="TopazSPower", short_models=("a62",)) + TOPAZSV = ProductInfo(nickname="TopazSV", short_models=("a26", "a27")) + # Ultron Series - "a73": RoborockProductNickname.ULTRONLITE, - "a85": RoborockProductNickname.ULTRONLITE, - "a94": RoborockProductNickname.ULTRONSC, - "a95": RoborockProductNickname.ULTRONSC, - "a124": RoborockProductNickname.ULTRONSE, - "a125": RoborockProductNickname.ULTRONSE, - "a139": RoborockProductNickname.ULTRONSE, - "a140": RoborockProductNickname.ULTRONSE, - "a68": RoborockProductNickname.ULTRONSPLUS, - "a69": RoborockProductNickname.ULTRONSPLUS, - "a70": RoborockProductNickname.ULTRONSPLUS, - "a50": RoborockProductNickname.ULTRON, - "a51": RoborockProductNickname.ULTRON, - "a72": RoborockProductNickname.ULTRONE, - "a84": RoborockProductNickname.ULTRONE, - "a96": RoborockProductNickname.ULTRONSV, - "a97": RoborockProductNickname.ULTRONSV, + ULTRON = ProductInfo(nickname="Ultron", short_models=("a50", "a51")) + ULTRONE = ProductInfo(nickname="UltronE", short_models=("a72", "a84")) + ULTRONLITE = ProductInfo(nickname="UltronLite", short_models=("a73", "a85")) + ULTRONSC = ProductInfo(nickname="UltronSC", short_models=("a94", "a95")) + ULTRONSE = ProductInfo(nickname="UltronSE", short_models=("a124", "a125", "a139", "a140")) + ULTRONSPLUS = ProductInfo(nickname="UltronSPlus", short_models=("a68", "a69", "a70")) + ULTRONSV = ProductInfo(nickname="UltronSV", short_models=("a96", "a97")) + # Verdelite Series - "a146": RoborockProductNickname.VERDELITE, - "a147": RoborockProductNickname.VERDELITE, - # Topaz Series - "a29": RoborockProductNickname.TOPAZS, - "a30": RoborockProductNickname.TOPAZS, - "a76": RoborockProductNickname.TOPAZS, - "a46": RoborockProductNickname.TOPAZSPLUS, - "a47": RoborockProductNickname.TOPAZSPLUS, - "a66": RoborockProductNickname.TOPAZSPLUS, - "a64": RoborockProductNickname.TOPAZSC, - "a65": RoborockProductNickname.TOPAZSC, - "a26": RoborockProductNickname.TOPAZSV, - "a27": RoborockProductNickname.TOPAZSV, - "a62": RoborockProductNickname.TOPAZSPOWER, - # Tanos Series - "a23": RoborockProductNickname.TANOSSPLUS, - "a24": RoborockProductNickname.TANOSSPLUS, - "a37": RoborockProductNickname.TANOSSLITE, - "a38": RoborockProductNickname.TANOSSLITE, - "a39": RoborockProductNickname.TANOSSC, - "a40": RoborockProductNickname.TANOSSC, - "a33": RoborockProductNickname.TANOSSE, - "a34": RoborockProductNickname.TANOSSE, - "a52": RoborockProductNickname.TANOSSMAX, - "t6": RoborockProductNickname.TANOS, - "s6": RoborockProductNickname.TANOS, - "t7": RoborockProductNickname.TANOSE, - "a11": RoborockProductNickname.TANOSE, - "t7p": RoborockProductNickname.TANOSV, - "a09": RoborockProductNickname.TANOSV, - "a10": RoborockProductNickname.TANOSV, - "a14": RoborockProductNickname.TANOSS, - "a15": RoborockProductNickname.TANOSS, - # Ruby Series - "t4": RoborockProductNickname.RUBYPLUS, - "s4": RoborockProductNickname.RUBYPLUS, - "p5": RoborockProductNickname.RUBYSC, - "a08": RoborockProductNickname.RUBYSC, - "a19": RoborockProductNickname.RUBYSE, - "p6": RoborockProductNickname.RUBYSLITE, - "s5e": RoborockProductNickname.RUBYSLITE, - "a05": RoborockProductNickname.RUBYSLITE, -} + VERDELITE = ProductInfo(nickname="Verdelite", short_models=("a146", "a147")) + + # Vivian Series + VIVIAN = ProductInfo(nickname="Vivian", short_models=("a134", "a135", "a155", "a156")) + VIVIANC = ProductInfo(nickname="VivianC", short_models=("a158", "a159")) + + +short_model_to_enum = {model: product for product in RoborockProduct for model in product.value.short_models} class RoborockStateCode(RoborockEnum): diff --git a/roborock/containers.py b/roborock/containers.py index 9acdb0f0..dfdc03f6 100644 --- a/roborock/containers.py +++ b/roborock/containers.py @@ -4,7 +4,7 @@ import json import logging import re -from dataclasses import asdict, dataclass, field +from dataclasses import asdict, dataclass, field, fields from datetime import timezone from enum import Enum, IntEnum from typing import Any, NamedTuple, get_args, get_origin @@ -354,104 +354,136 @@ class DeviceFeatures(RoborockBase): """Represents the features supported by a Roborock device.""" # Features derived from robot_new_features - is_map_carpet_add_support: bool - is_show_clean_finish_reason_supported: bool - is_resegment_supported: bool - is_video_monitor_supported: bool - is_any_state_transit_goto_supported: bool - is_fw_filter_obstacle_supported: bool - is_video_settings_supported: bool - is_ignore_unknown_map_object_supported: bool - is_set_child_supported: bool - is_carpet_supported: bool - is_mop_path_supported: bool - is_multi_map_segment_timer_supported: bool - is_custom_water_box_distance_supported: bool - is_wash_then_charge_cmd_supported: bool - is_room_name_supported: bool - is_current_map_restore_enabled: bool - is_photo_upload_supported: bool - is_shake_mop_set_supported: bool - is_map_beautify_internal_debug_supported: bool - is_new_data_for_clean_history_supported: bool - is_new_data_for_clean_history_detail_supported: bool - is_flow_led_setting_supported: bool - is_dust_collection_setting_supported: bool - is_rpc_retry_supported: bool - is_avoid_collision_supported: bool - is_support_set_switch_map_mode_supported: bool - is_support_smart_scene_supported: bool - is_support_floor_edit_supported: bool - is_support_furniture_supported: bool - is_support_room_tag_supported: bool - is_support_quick_map_builder_supported: bool - is_support_smart_global_clean_with_custom_mode_supported: bool - is_record_allowed: bool - is_careful_slow_mop_supported: bool - is_egg_mode_supported: bool - is_carpet_show_on_map_supported: bool - is_supported_valley_electricity_supported: bool - is_unsave_map_reason_supported: bool - is_supported_drying_supported: bool - is_supported_download_test_voice_supported: bool - is_support_backup_map_supported: bool - is_support_custom_mode_in_cleaning_supported: bool - is_support_remote_control_in_call_supported: bool + is_show_clean_finish_reason_supported: bool = field(metadata={"robot_new_features": 1}) + is_resegment_supported: bool = field(metadata={"robot_new_features": 4}) + is_video_monitor_supported: bool = field(metadata={"robot_new_features": 8}) + is_any_state_transit_goto_supported: bool = field(metadata={"robot_new_features": 16}) + is_fw_filter_obstacle_supported: bool = field(metadata={"robot_new_features": 32}) + is_video_settings_supported: bool = field(metadata={"robot_new_features": 64}) + is_ignore_unknown_map_object_supported: bool = field(metadata={"robot_new_features": 128}) + is_set_child_supported: bool = field(metadata={"robot_new_features": 256}) + is_carpet_supported: bool = field(metadata={"robot_new_features": 512}) + is_record_allowed: bool = field(metadata={"robot_new_features": 1024}) + is_mop_path_supported: bool = field(metadata={"robot_new_features": 2048}) + is_current_map_restore_enabled: bool = field(metadata={"robot_new_features": 8192}) + is_room_name_supported: bool = field(metadata={"robot_new_features": 16384}) + is_photo_upload_supported: bool = field(metadata={"robot_new_features": 65536}) + is_shake_mop_set_supported: bool = field(metadata={"robot_new_features": 262144}) + is_map_beautify_internal_debug_supported: bool = field(metadata={"robot_new_features": 2097152}) + is_new_data_for_clean_history_supported: bool = field(metadata={"robot_new_features": 4194304}) + is_new_data_for_clean_history_detail_supported: bool = field(metadata={"robot_new_features": 8388608}) + is_flow_led_setting_supported: bool = field(metadata={"robot_new_features": 16777216}) + is_dust_collection_setting_supported: bool = field(metadata={"robot_new_features": 33554432}) + is_rpc_retry_supported: bool = field(metadata={"robot_new_features": 67108864}) + is_avoid_collision_supported: bool = field(metadata={"robot_new_features": 134217728}) + is_support_set_switch_map_mode_supported: bool = field(metadata={"robot_new_features": 268435456}) + is_map_carpet_add_support: bool = field(metadata={"robot_new_features": 1073741824}) + is_custom_water_box_distance_supported: bool = field(metadata={"robot_new_features": 2147483648}) # Features derived from unhexed_feature_info - is_support_set_volume_in_call: bool - is_support_clean_estimate: bool - is_support_custom_dnd: bool - is_carpet_deep_clean_supported: bool - is_support_stuck_zone: bool - is_support_custom_door_sill: bool - is_wifi_manage_supported: bool - is_clean_route_fast_mode_supported: bool - is_support_cliff_zone: bool - is_support_smart_door_sill: bool - is_support_floor_direction: bool - is_back_charge_auto_wash_supported: bool - is_super_deep_wash_supported: bool - is_ces2022_supported: bool - is_dss_believable_supported: bool - is_main_brush_up_down_supported: bool - is_goto_pure_clean_path_supported: bool - is_water_up_down_drain_supported: bool - is_setting_carpet_first_supported: bool - is_clean_route_deep_slow_plus_supported: bool - is_left_water_drain_supported: bool - is_clean_count_setting_supported: bool - is_corner_clean_mode_supported: bool + is_support_smart_scene_supported: bool = field(metadata={"upper_32_bits": 1}) + is_support_floor_edit_supported: bool = field(metadata={"upper_32_bits": 3}) + is_support_furniture_supported: bool = field(metadata={"upper_32_bits": 4}) + is_wash_then_charge_cmd_supported: bool = field(metadata={"upper_32_bits": 5}) + is_support_room_tag_supported: bool = field(metadata={"upper_32_bits": 6}) + is_support_quick_map_builder_supported: bool = field(metadata={"upper_32_bits": 7}) + is_support_smart_global_clean_with_custom_mode_supported: bool = field(metadata={"upper_32_bits": 8}) + is_careful_slow_mop_supported: bool = field(metadata={"upper_32_bits": 9}) + is_egg_mode_supported: bool = field(metadata={"upper_32_bits": 10}) + is_carpet_show_on_map_supported: bool = field(metadata={"upper_32_bits": 12}) + is_supported_valley_electricity_supported: bool = field(metadata={"upper_32_bits": 13}) + is_unsave_map_reason_supported: bool = field(metadata={"upper_32_bits": 14}) + is_supported_download_test_voice_supported: bool = field(metadata={"upper_32_bits": 16}) + is_support_backup_map_supported: bool = field(metadata={"upper_32_bits": 17}) + is_support_custom_mode_in_cleaning_supported: bool = field(metadata={"upper_32_bits": 18}) + is_support_remote_control_in_call_supported: bool = field(metadata={"upper_32_bits": 19}) + + is_support_set_volume_in_call: bool = field(metadata={"unhexed_feature_info": 1}) + is_support_clean_estimate: bool = field(metadata={"unhexed_feature_info": 2}) + is_support_custom_dnd: bool = field(metadata={"unhexed_feature_info": 4}) + is_carpet_deep_clean_supported: bool = field(metadata={"unhexed_feature_info": 8}) + is_support_stuck_zone: bool = field(metadata={"unhexed_feature_info": 16}) + is_support_custom_door_sill: bool = field(metadata={"unhexed_feature_info": 32}) + is_wifi_manage_supported: bool = field(metadata={"unhexed_feature_info": 128}) + is_clean_route_fast_mode_supported: bool = field(metadata={"unhexed_feature_info": 256}) + is_support_cliff_zone: bool = field(metadata={"unhexed_feature_info": 512}) + is_support_smart_door_sill: bool = field(metadata={"unhexed_feature_info": 1024}) + is_support_floor_direction: bool = field(metadata={"unhexed_feature_info": 2048}) + is_back_charge_auto_wash_supported: bool = field(metadata={"unhexed_feature_info": 4096}) + is_super_deep_wash_supported: bool = field(metadata={"unhexed_feature_info": 32768}) + is_ces2022_supported: bool = field(metadata={"unhexed_feature_info": 65536}) + is_dss_believable_supported: bool = field(metadata={"unhexed_feature_info": 131072}) + is_main_brush_up_down_supported: bool = field(metadata={"unhexed_feature_info": 262144}) + is_goto_pure_clean_path_supported: bool = field(metadata={"unhexed_feature_info": 524288}) + is_water_up_down_drain_supported: bool = field(metadata={"unhexed_feature_info": 1048576}) + is_setting_carpet_first_supported: bool = field(metadata={"unhexed_feature_info": 8388608}) + is_clean_route_deep_slow_plus_supported: bool = field(metadata={"unhexed_feature_info": 16777216}) + is_left_water_drain_supported: bool = field(metadata={"unhexed_feature_info": 134217728}) + is_clean_count_setting_supported: bool = field(metadata={"unhexed_feature_info": 1073741824}) + is_corner_clean_mode_supported: bool = field(metadata={"unhexed_feature_info": 2147483648}) # --- Features from new_feature_info_str --- - is_two_key_real_time_video_supported: bool - is_two_key_rtv_in_charging_supported: bool - is_dirty_replenish_clean_supported: bool - is_avoid_collision_mode_str_supported: bool - is_voice_control_str_supported: bool - is_new_endpoint_supported: bool - is_corner_mop_strech_supported: bool - is_hot_wash_towel_supported: bool - is_floor_dir_clean_any_time_supported: bool - is_pet_supplies_deep_clean_supported: bool - is_mop_shake_water_max_supported: bool - is_custom_clean_mode_count_supported: bool - is_exact_custom_mode_supported: bool - is_carpet_custom_clean_supported: bool - is_pet_snapshot_supported: bool - is_new_ai_recognition_supported: bool - is_auto_collection_2_supported: bool - is_right_brush_stretch_supported: bool - is_smart_clean_mode_set_supported: bool - is_dirty_object_detect_supported: bool - is_no_need_carpet_press_set_supported: bool - is_voice_control_led_supported: bool - is_water_leak_check_supported: bool - is_min_battery_15_to_clean_task_supported: bool - is_gap_deep_clean_supported: bool - is_object_detect_check_supported: bool - is_identify_room_supported: bool - is_matter_supported: bool + is_two_key_real_time_video_supported: bool = field( + metadata={"new_feature_str_bit": NewFeatureStrBit.TWO_KEY_REAL_TIME_VIDEO} + ) + is_two_key_rtv_in_charging_supported: bool = field( + metadata={"new_feature_str_bit": NewFeatureStrBit.TWO_KEY_RTV_IN_CHARGING} + ) + is_dirty_replenish_clean_supported: bool = field( + metadata={"new_feature_str_bit": NewFeatureStrBit.DIRTY_REPLENISH_CLEAN} + ) + is_avoid_collision_mode_str_supported: bool = field( + metadata={"new_feature_str_bit": NewFeatureStrBit.AVOID_COLLISION_MODE} + ) + is_voice_control_str_supported: bool = field(metadata={"new_feature_str_bit": NewFeatureStrBit.VOICE_CONTROL}) + is_new_endpoint_supported: bool = field(metadata={"new_feature_str_bit": NewFeatureStrBit.NEW_ENDPOINT}) + is_corner_mop_strech_supported: bool = field(metadata={"new_feature_str_bit": NewFeatureStrBit.CORNER_MOP_STRECH}) + is_hot_wash_towel_supported: bool = field(metadata={"new_feature_str_bit": NewFeatureStrBit.HOT_WASH_TOWEL}) + is_floor_dir_clean_any_time_supported: bool = field( + metadata={"new_feature_str_bit": NewFeatureStrBit.FLOOR_DIR_CLEAN_ANY_TIME} + ) + is_pet_supplies_deep_clean_supported: bool = field( + metadata={"new_feature_str_bit": NewFeatureStrBit.PET_SUPPLIES_DEEP_CLEAN} + ) + is_mop_shake_water_max_supported: bool = field( + metadata={"new_feature_str_bit": NewFeatureStrBit.MOP_SHAKE_WATER_MAX} + ) + is_exact_custom_mode_supported: bool = field(metadata={"new_feature_str_bit": NewFeatureStrBit.EXACT_CUSTOM_MODE}) + is_carpet_custom_clean_supported: bool = field( + metadata={"new_feature_str_bit": NewFeatureStrBit.CARPET_CUSTOM_CLEAN} + ) + is_pet_snapshot_supported: bool = field(metadata={"new_feature_str_bit": NewFeatureStrBit.PET_SNAPSHOT}) + is_custom_clean_mode_count_supported: bool = field( + metadata={"new_feature_str_bit": NewFeatureStrBit.CUSTOM_CLEAN_MODE_COUNT} + ) + is_new_ai_recognition_supported: bool = field(metadata={"new_feature_str_bit": NewFeatureStrBit.NEW_AI_RECOGNITION}) + is_auto_collection_2_supported: bool = field(metadata={"new_feature_str_bit": NewFeatureStrBit.AUTO_COLLECTION_2}) + is_right_brush_stretch_supported: bool = field( + metadata={"new_feature_str_bit": NewFeatureStrBit.RIGHT_BRUSH_STRETCH} + ) + is_smart_clean_mode_set_supported: bool = field( + metadata={"new_feature_str_bit": NewFeatureStrBit.SMART_CLEAN_MODE_SET} + ) + is_dirty_object_detect_supported: bool = field( + metadata={"new_feature_str_bit": NewFeatureStrBit.DIRTY_OBJECT_DETECT} + ) + is_no_need_carpet_press_set_supported: bool = field( + metadata={"new_feature_str_bit": NewFeatureStrBit.NO_NEED_CARPET_PRESS_SET} + ) + is_voice_control_led_supported: bool = field(metadata={"new_feature_str_bit": NewFeatureStrBit.VOICE_CONTROL_LED}) + is_water_leak_check_supported: bool = field(metadata={"new_feature_str_bit": NewFeatureStrBit.WATER_LEAK_CHECK}) + is_min_battery_15_to_clean_task_supported: bool = field( + metadata={"new_feature_str_bit": NewFeatureStrBit.MIN_BATTERY_15_TO_CLEAN_TASK} + ) + is_gap_deep_clean_supported: bool = field(metadata={"new_feature_str_bit": NewFeatureStrBit.GAP_DEEP_CLEAN}) + is_object_detect_check_supported: bool = field( + metadata={"new_feature_str_bit": NewFeatureStrBit.OBJECT_DETECT_CHECK} + ) + is_identify_room_supported: bool = field(metadata={"new_feature_str_bit": NewFeatureStrBit.IDENTIFY_ROOM}) + is_matter_supported: bool = field(metadata={"new_feature_str_bit": NewFeatureStrBit.MATTER}) + + # is_multi_map_segment_timer_supported: bool = field(default=False) + # is_supported_drying_supported: bool = field(default=False) @classmethod def _is_new_feature_str_support(cls, o: int, new_feature_info_str: str) -> bool: @@ -461,15 +493,11 @@ def _is_new_feature_str_support(cls, o: int, new_feature_info_str: str) -> bool: try: l = o % 4 target_index = -((o // 4) + 1) - p = new_feature_info_str[target_index] - hex_char_value = int(p, 16) - is_set = (hex_char_value >> l) & 1 - return bool(is_set) - except Exception: + except (IndexError, ValueError): return False @classmethod @@ -477,158 +505,29 @@ def from_feature_flags( cls, robot_new_features: int, new_feature_set: str, product_nickname: RoborockProductNickname ) -> DeviceFeatures: """Creates a DeviceFeatures instance from raw feature flags.""" - unhexed_feature_info = int(new_feature_set[-8:], 16) - + unhexed_feature_info = int(new_feature_set[-8:], 16) if new_feature_set and len(new_feature_set) >= 8 else 0 upper_32_bits = robot_new_features // (2**32) - return cls( - is_map_carpet_add_support=bool(1073741824 & robot_new_features), - is_show_clean_finish_reason_supported=bool(1 & robot_new_features), - is_resegment_supported=bool(4 & robot_new_features), - is_video_monitor_supported=bool(8 & robot_new_features), - is_any_state_transit_goto_supported=bool(16 & robot_new_features), - is_fw_filter_obstacle_supported=bool(32 & robot_new_features), - is_video_settings_supported=bool(64 & robot_new_features), - is_ignore_unknown_map_object_supported=bool(128 & robot_new_features), - is_set_child_supported=bool(256 & robot_new_features), - is_carpet_supported=bool(512 & robot_new_features), - is_mop_path_supported=bool(2048 & robot_new_features), - is_multi_map_segment_timer_supported=False, # TODO - is_custom_water_box_distance_supported=bool(2147483648 & robot_new_features), - is_wash_then_charge_cmd_supported=bool(robot_new_features and ((upper_32_bits >> 5) & 1)), - is_room_name_supported=bool(16384 & robot_new_features), - is_current_map_restore_enabled=bool(8192 & robot_new_features), - is_photo_upload_supported=bool(65536 & robot_new_features), - is_shake_mop_set_supported=bool(262144 & robot_new_features), - is_map_beautify_internal_debug_supported=bool(2097152 & robot_new_features), - is_new_data_for_clean_history_supported=bool(4194304 & robot_new_features), - is_new_data_for_clean_history_detail_supported=bool(8388608 & robot_new_features), - is_flow_led_setting_supported=bool(16777216 & robot_new_features), - is_dust_collection_setting_supported=bool(33554432 & robot_new_features), - is_rpc_retry_supported=bool(67108864 & robot_new_features), - is_avoid_collision_supported=bool(134217728 & robot_new_features), - is_support_set_switch_map_mode_supported=bool(268435456 & robot_new_features), - is_support_smart_scene_supported=bool(robot_new_features and (upper_32_bits & 2)), - is_support_floor_edit_supported=bool(robot_new_features and (upper_32_bits & 8)), - is_support_furniture_supported=bool(robot_new_features and ((upper_32_bits >> 4) & 1)), - is_support_room_tag_supported=bool(robot_new_features and ((upper_32_bits >> 6) & 1)), - is_support_quick_map_builder_supported=bool(robot_new_features and ((upper_32_bits >> 7) & 1)), - is_support_smart_global_clean_with_custom_mode_supported=bool( - robot_new_features and ((upper_32_bits >> 8) & 1) - ), - is_record_allowed=bool(1024 & robot_new_features), - is_careful_slow_mop_supported=bool(robot_new_features and ((upper_32_bits >> 9) & 1)), - is_egg_mode_supported=bool(robot_new_features and ((upper_32_bits >> 10) & 1)), - is_carpet_show_on_map_supported=bool(robot_new_features and ((upper_32_bits >> 12) & 1)), - is_supported_valley_electricity_supported=bool(robot_new_features and ((upper_32_bits >> 13) & 1)), - is_unsave_map_reason_supported=bool(robot_new_features and ((upper_32_bits >> 14) & 1)), - is_supported_drying_supported=False, # TODO - is_supported_download_test_voice_supported=bool(robot_new_features and ((upper_32_bits >> 16) & 1)), - is_support_backup_map_supported=bool(robot_new_features and ((upper_32_bits >> 17) & 1)), - is_support_custom_mode_in_cleaning_supported=bool(robot_new_features and ((upper_32_bits >> 18) & 1)), - is_support_remote_control_in_call_supported=bool(robot_new_features and ((upper_32_bits >> 19) & 1)), - # Features from unhexed_feature_info - is_support_set_volume_in_call=bool(1 & unhexed_feature_info), - is_support_clean_estimate=bool(2 & unhexed_feature_info), - is_support_custom_dnd=bool(4 & unhexed_feature_info), - is_carpet_deep_clean_supported=bool(8 & unhexed_feature_info), - is_support_stuck_zone=bool(16 & unhexed_feature_info), - is_support_custom_door_sill=bool(32 & unhexed_feature_info), - is_wifi_manage_supported=bool(128 & unhexed_feature_info), - is_clean_route_fast_mode_supported=bool(256 & unhexed_feature_info), - is_support_cliff_zone=bool(512 & unhexed_feature_info), - is_support_smart_door_sill=bool(1024 & unhexed_feature_info), - is_support_floor_direction=bool(2048 & unhexed_feature_info), - is_back_charge_auto_wash_supported=bool(4096 & unhexed_feature_info), - is_super_deep_wash_supported=bool(32768 & unhexed_feature_info), - is_ces2022_supported=bool(65536 & unhexed_feature_info), - is_dss_believable_supported=bool(131072 & unhexed_feature_info), - is_main_brush_up_down_supported=bool(262144 & unhexed_feature_info), - is_goto_pure_clean_path_supported=bool(524288 & unhexed_feature_info), - is_water_up_down_drain_supported=bool(1048576 & unhexed_feature_info), - is_setting_carpet_first_supported=bool(8388608 & unhexed_feature_info), - is_clean_route_deep_slow_plus_supported=bool(16777216 & unhexed_feature_info), - is_left_water_drain_supported=bool(134217728 & unhexed_feature_info), - is_clean_count_setting_supported=bool(1073741824 & unhexed_feature_info), - is_corner_clean_mode_supported=bool(2147483648 & unhexed_feature_info), - # Features from is_new_feature_str_support - is_two_key_real_time_video_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.TWO_KEY_REAL_TIME_VIDEO, new_feature_set - ), - is_two_key_rtv_in_charging_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.TWO_KEY_RTV_IN_CHARGING, new_feature_set - ), - is_dirty_replenish_clean_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.DIRTY_REPLENISH_CLEAN, new_feature_set - ), - is_avoid_collision_mode_str_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.AVOID_COLLISION_MODE, new_feature_set - ), - is_voice_control_str_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.VOICE_CONTROL, new_feature_set - ), - is_new_endpoint_supported=cls._is_new_feature_str_support(NewFeatureStrBit.NEW_ENDPOINT, new_feature_set), - is_corner_mop_strech_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.CORNER_MOP_STRECH, new_feature_set - ), - is_hot_wash_towel_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.HOT_WASH_TOWEL, new_feature_set - ), - is_floor_dir_clean_any_time_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.FLOOR_DIR_CLEAN_ANY_TIME, new_feature_set - ), - is_pet_supplies_deep_clean_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.PET_SUPPLIES_DEEP_CLEAN, new_feature_set - ), - is_mop_shake_water_max_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.MOP_SHAKE_WATER_MAX, new_feature_set - ), - is_custom_clean_mode_count_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.CUSTOM_CLEAN_MODE_COUNT, new_feature_set - ), - is_exact_custom_mode_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.EXACT_CUSTOM_MODE, new_feature_set - ), - is_carpet_custom_clean_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.CARPET_CUSTOM_CLEAN, new_feature_set - ), - is_pet_snapshot_supported=cls._is_new_feature_str_support(NewFeatureStrBit.PET_SNAPSHOT, new_feature_set), - is_new_ai_recognition_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.NEW_AI_RECOGNITION, new_feature_set - ), - is_auto_collection_2_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.AUTO_COLLECTION_2, new_feature_set - ), - is_right_brush_stretch_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.RIGHT_BRUSH_STRETCH, new_feature_set - ), - is_smart_clean_mode_set_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.SMART_CLEAN_MODE_SET, new_feature_set - ), - is_dirty_object_detect_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.DIRTY_OBJECT_DETECT, new_feature_set - ), - is_no_need_carpet_press_set_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.NO_NEED_CARPET_PRESS_SET, new_feature_set - ), - is_voice_control_led_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.VOICE_CONTROL_LED, new_feature_set - ), - is_water_leak_check_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.WATER_LEAK_CHECK, new_feature_set - ), - is_min_battery_15_to_clean_task_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.MIN_BATTERY_15_TO_CLEAN_TASK, new_feature_set - ), - is_gap_deep_clean_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.GAP_DEEP_CLEAN, new_feature_set - ), - is_object_detect_check_supported=cls._is_new_feature_str_support( - NewFeatureStrBit.OBJECT_DETECT_CHECK, new_feature_set - ), - is_identify_room_supported=cls._is_new_feature_str_support(NewFeatureStrBit.IDENTIFY_ROOM, new_feature_set), - is_matter_supported=cls._is_new_feature_str_support(NewFeatureStrBit.MATTER, new_feature_set), - ) + kwargs: dict[str, Any] = {} + + for f in fields(cls): + if not f.metadata: + continue + + if "robot_new_features" in f.metadata: + mask = f.metadata["robot_new_features"] + kwargs[f.name] = bool(mask & robot_new_features) + elif "upper_32_bits" in f.metadata: + bit_index = f.metadata["upper_32_bits"] + kwargs[f.name] = bool(robot_new_features and ((upper_32_bits >> bit_index) & 1)) + elif "unhexed_feature_info" in f.metadata: + mask = f.metadata["unhexed_feature_info"] + kwargs[f.name] = bool(mask & unhexed_feature_info) + elif "new_feature_str_bit" in f.metadata: + bit = f.metadata["new_feature_str_bit"] + kwargs[f.name] = cls._is_new_feature_str_support(bit, new_feature_set) + + return cls(**kwargs) @dataclass diff --git a/roborock/device_trait.py b/roborock/device_trait.py index f8837fd0..a767072a 100644 --- a/roborock/device_trait.py +++ b/roborock/device_trait.py @@ -3,7 +3,7 @@ from dataclasses import dataclass from . import RoborockCommand -from .containers import DeviceFeatures, RoborockBase +from .containers import DeviceFeatures @dataclass @@ -12,7 +12,6 @@ class DeviceTrait(ABC): def __init__(self, send_command: Callable[..., Awaitable[None]]): self.send_command = send_command - self.status: RoborockBase | None = None self.subscriptions = [] @classmethod @@ -21,11 +20,11 @@ def supported(cls, features: DeviceFeatures) -> bool: raise NotImplementedError @abstractmethod - def from_dict(cls, data: dict) -> bool: + def update(cls, data: dict) -> bool: raise NotImplementedError def on_message(self, data: dict) -> None: - self.status = self.from_dict(data) + self.status = self.update(data) for callback in self.subscriptions: callback(self.status) diff --git a/roborock/roborock_device.py b/roborock/roborock_device.py index e29b5916..7b1f5308 100644 --- a/roborock/roborock_device.py +++ b/roborock/roborock_device.py @@ -30,7 +30,6 @@ def __init__(self, user_data: UserData, device_info: DeviceData): rriot = user_data.rriot hashed_user = md5hex(rriot.u + ":" + rriot.k)[2:10] url = urlparse(rriot.r.m) - mqtt_password = rriot.s self._local_endpoint = "abc" self._nonce = secrets.token_bytes(16) diff --git a/roborock/roborock_message.py b/roborock/roborock_message.py index 1e8535b0..5f46c497 100644 --- a/roborock/roborock_message.py +++ b/roborock/roborock_message.py @@ -4,6 +4,7 @@ import math import time from dataclasses import dataclass, field +from functools import cache from roborock import RoborockEnum from roborock.util import get_next_int @@ -163,15 +164,13 @@ class RoborockMessage: message_retry: MessageRetry | None = None _parsed_payload: dict | None = None + @cache def get_payload(self) -> dict | None: - if self.payload and not self._parsed_payload: - self._parsed_payload = json.loads(self.payload.decode()) - return self._parsed_payload + return json.loads(self.payload.decode()) def get_request_id(self) -> int | None: - payload = self.get_payload() - if payload: - for data_point_number, data_point in self._parsed_payload.get("dps").items(): + if payload := self.get_payload(): + for data_point_number, data_point in payload.get("dps").items(): if data_point_number in ["101", "102"]: data_point_response = json.loads(data_point) return data_point_response.get("id") @@ -186,9 +185,8 @@ def get_method(self) -> str | None: if self.message_retry: return self.message_retry.method protocol = self.protocol - payload = self.get_payload() - if payload and protocol in [4, 5, 101, 102]: - for data_point_number, data_point in self._parsed_payload.get("dps").items(): + if payload := self.get_payload() and protocol in [4, 5, 101, 102]: + for data_point_number, data_point in payload.get("dps").items(): if data_point_number in ["101", "102"]: data_point_response = json.loads(data_point) return data_point_response.get("method") @@ -196,9 +194,8 @@ def get_method(self) -> str | None: def get_params(self) -> list | dict | None: protocol = self.protocol - payload = self.get_payload() - if payload and protocol in [4, 101, 102]: - for data_point_number, data_point in self._parsed_payload.get("dps").items(): + if payload := self.get_payload() and protocol in [4, 101, 102]: + for data_point_number, data_point in payload.get("dps").items(): if data_point_number in ["101", "102"]: data_point_response = json.loads(data_point) return data_point_response.get("params")