feat: Add zeo support and fix some a01 weirdness (#200)

Lash-L · web-flow · commit e825ff581151 · 2024-04-11T10:46:01.000-04:00
* major: add A01

* chore: add init

* chore: fix commitlint?

* chore: fix commitlint

* chore: fix commitlint

* chore: change refactor to be major tag

* refactor: add A01

* feat: add a01

BREAKING CHANGE: You must now specify what version api you want to use with clients.

* feat: add initial zeo support

* fix: fix A01 support

* fix: allow messages to fail

* fix: lint

* feat: add more zeo things
diff --git a/roborock/api.py b/roborock/api.py
@@ -12,9 +12,6 @@
 
 from .containers import (
     DeviceData,
-    ModelStatus,
-    S7MaxVStatus,
-    Status,
 )
 from .exceptions import (
     RoborockTimeout,
@@ -48,16 +45,10 @@ def __init__(self, endpoint: str, device_info: DeviceData, queue_timeout: int =
         self._logger = RoborockLoggerAdapter(device_info.device.name, _LOGGER)
         self.is_available: bool = True
         self.queue_timeout = queue_timeout
-        self._status_type: type[Status] = ModelStatus.get(self.device_info.model, S7MaxVStatus)
 
     def __del__(self) -> None:
         self.release()
 
-    @property
-    def status_type(self) -> type[Status]:
-        """Gets the status type for this device"""
-        return self._status_type
-
     def release(self):
         self.sync_disconnect()
 
diff --git a/roborock/code_mappings.py b/roborock/code_mappings.py
@@ -450,3 +450,138 @@ class DyadError(RoborockEnum):
     dirty_charging_contacts = 10007  # Disconnection between the device and dock. Wipe charging contacts.
     low_battery = 20017  # Low battery level. Charge before starting self-cleaning.
     battery_under_10 = 20018  # Charge until the battery level exceeds 10% before manually starting self-cleaning.
+
+
+class ZeoMode(RoborockEnum):
+    wash = 1
+    wash_and_dry = 2
+    dry = 3
+
+
+class ZeoState(RoborockEnum):
+    standby = 1
+    weighing = 2
+    soaking = 3
+    washing = 4
+    rinsing = 5
+    spinning = 6
+    drying = 7
+    cooling = 8
+    under_delay_start = 9
+    done = 10
+
+
+class ZeoProgram(RoborockEnum):
+    standard = 1
+    quick = 2
+    sanitize = 3
+    wool = 4
+    air_refresh = 5
+    custom = 6
+    bedding = 7
+    down = 8
+    silk = 9
+    rinse_and_spin = 10
+    spin = 11
+    down_clean = 12
+    baby_care = 13
+    anti_allergen = 14
+    sportswear = 15
+    night = 16
+    new_clothes = 17
+    shirts = 18
+    synthetics = 19
+    underwear = 20
+    gentle = 21
+    intensive = 22
+    cotton_linen = 23
+    season = 24
+    warming = 25
+    bra = 26
+    panties = 27
+    boiling_wash = 28
+    socks = 30
+    towels = 31
+    anti_mite = 32
+    exo_40_60 = 33
+    twenty_c = 34
+    t_shirts = 35
+    stain_removal = 36
+
+
+class ZeoSoak(RoborockEnum):
+    normal = 0
+    low = 1
+    medium = 2
+    high = 3
+    max = 4
+
+
+class ZeoTemperature(RoborockEnum):
+    normal = 1
+    low = 2
+    medium = 3
+    high = 4
+    max = 5
+    twenty_c = 6
+
+
+class ZeoRinse(RoborockEnum):
+    none = 0
+    min = 1
+    low = 2
+    mid = 3
+    high = 4
+    max = 5
+
+
+class ZeoSpin(RoborockEnum):
+    none = 1
+    very_low = 2
+    low = 3
+    mid = 4
+    high = 5
+    very_high = 6
+    max = 7
+
+
+class ZeoDryingMode(RoborockEnum):
+    none = 0
+    quick = 1
+    iron = 2
+    store = 3
+
+
+class ZeoDetergentType(RoborockEnum):
+    empty = 0
+    low = 1
+    medium = 2
+    high = 3
+
+
+class ZeoSoftenerType(RoborockEnum):
+    empty = 0
+    low = 1
+    medium = 2
+    high = 3
+
+
+class ZeoError(RoborockEnum):
+    none = 0
+    refill_error = 1
+    drain_error = 2
+    door_lock_error = 3
+    water_level_error = 4
+    inverter_error = 5
+    heating_error = 6
+    temperature_error = 7
+    communication_error = 10
+    drying_error = 11
+    drying_error_e_12 = 12
+    drying_error_e_13 = 13
+    drying_error_e_14 = 14
+    drying_error_e_15 = 15
+    drying_error_e_16 = 16
+    drying_error_water_flow = 17  # Check for normal water flow
+    drying_error_restart = 18  # Restart the washer and try again
+    spin_error = 19  # re-arrange clothes
diff --git a/roborock/protocol.py b/roborock/protocol.py
@@ -36,7 +36,6 @@
 _LOGGER = logging.getLogger(__name__)
 SALT = b"TXdfu$jyZ#TZHsg4"
 A01_HASH = "726f626f726f636b2d67a6d6da"
-A01_AES_DECIPHER = "ELSYN0wTI4AUm7C4"
 BROADCAST_TOKEN = b"qWKYcdQWrbm9hPqe"
 AP_CONFIG = 1
 SOCK_DISCOVERY = 2
@@ -208,7 +207,7 @@ def _encode(self, obj, context, _):
         """
         if context.version == b"A01":
             iv = md5hex(format(context.random, "08x") + A01_HASH)[8:24]
-            decipher = AES.new(bytes(A01_AES_DECIPHER, "utf-8"), AES.MODE_CBC, bytes(iv, "utf-8"))
+            decipher = AES.new(bytes(context.search("local_key"), "utf-8"), AES.MODE_CBC, bytes(iv, "utf-8"))
             f = decipher.encrypt(obj)
             return f
         token = self.token_func(context)
@@ -219,7 +218,7 @@ def _decode(self, obj, context, _):
         """Decrypts the given payload with the token stored in the context."""
         if context.version == b"A01":
             iv = md5hex(format(context.random, "08x") + A01_HASH)[8:24]
-            decipher = AES.new(bytes(A01_AES_DECIPHER, "utf-8"), AES.MODE_CBC, bytes(iv, "utf-8"))
+            decipher = AES.new(bytes(context.search("local_key"), "utf-8"), AES.MODE_CBC, bytes(iv, "utf-8"))
             f = decipher.decrypt(obj)
             return f
         token = self.token_func(context)
diff --git a/roborock/roborock_message.py b/roborock/roborock_message.py
@@ -87,6 +87,48 @@ class RoborockDyadDataProtocol(RoborockEnum):
     RPC_RESPONSE = 10102
 
 
+class RoborockZeoProtocol(RoborockEnum):
+    START = 200  # rw
+    PAUSE = 201  # rw
+    SHUTDOWN = 202  # rw
+    STATE = 203  # ro
+    MODE = 204  # rw
+    PROGRAM = 205  # rw
+    CHILD_LOCK = 206  # rw
+    TEMP = 207  # rw
+    RINSE_TIMES = 208  # rw
+    SPIN_LEVEL = 209  # rw
+    DRYING_MODE = 210  # rw
+    DETERGENT_SET = 211  # rw
+    SOFTENER_SET = 212  # rw
+    DETERGENT_TYPE = 213  # rw
+    SOFTENER_TYPE = 214  # rw
+    COUNTDOWN = 217  # rw
+    WASHING_LEFT = 218  # ro
+    DOORLOCK_STATE = 219  # ro
+    ERROR = 220  # ro
+    CUSTOM_PARAM_SAVE = 221  # rw
+    CUSTOM_PARAM_GET = 222  # ro
+    SOUND_SET = 223  # rw
+    TIMES_AFTER_CLEAN = 224  # ro
+    DEFAULT_SETTING = 225  # rw
+    DETERGENT_EMPTY = 226  # ro
+    SOFTENER_EMPTY = 227  # ro
+    LIGHT_SETTING = 229  # rw
+    DETERGENT_VOLUME = 230  # rw
+    SOFTENER_VOLUME = 231  # rw
+    APP_AUTHORIZATION = 232  # rw
+    ID_QUERY = 10000
+    F_C = 10001
+    SND_STATE = 10004
+    PRODUCT_INFO = 10005
+    PRIVACY_INFO = 10006
+    OTA_NFO = 10007
+    WASHING_LOG = 10008
+    RPC_REQ = 10101
+    RPC_RESp = 10102
+
+
 ROBOROCK_DATA_STATUS_PROTOCOL = [
     RoborockDataProtocol.ERROR_CODE,
     RoborockDataProtocol.STATE,
diff --git a/roborock/version_a01_apis/roborock_client_a01.py b/roborock/version_a01_apis/roborock_client_a01.py
@@ -19,59 +19,92 @@
     DyadWarmLevel,
     DyadWaterLevel,
     RoborockDyadStateCode,
+    ZeoDetergentType,
+    ZeoDryingMode,
+    ZeoError,
+    ZeoMode,
+    ZeoProgram,
+    ZeoRinse,
+    ZeoSoftenerType,
+    ZeoSpin,
+    ZeoState,
+    ZeoTemperature,
 )
-from roborock.containers import DyadProductInfo, DyadSndState
+from roborock.containers import DyadProductInfo, DyadSndState, RoborockCategory
 from roborock.roborock_message import (
     RoborockDyadDataProtocol,
     RoborockMessage,
     RoborockMessageProtocol,
+    RoborockZeoProtocol,
 )
 
 
 @dataclasses.dataclass
-class DyadProtocolCacheEntry:
+class A01ProtocolCacheEntry:
     post_process_fn: Callable
     value: typing.Any | None = None
 
 
 # Right now this cache is not active, it was too much complexity for the initial addition of dyad.
 protocol_entries = {
-    RoborockDyadDataProtocol.STATUS: DyadProtocolCacheEntry(lambda val: RoborockDyadStateCode(val).name),
-    RoborockDyadDataProtocol.SELF_CLEAN_MODE: DyadProtocolCacheEntry(lambda val: DyadSelfCleanMode(val).name),
-    RoborockDyadDataProtocol.SELF_CLEAN_LEVEL: DyadProtocolCacheEntry(lambda val: DyadSelfCleanLevel(val).name),
-    RoborockDyadDataProtocol.WARM_LEVEL: DyadProtocolCacheEntry(lambda val: DyadWarmLevel(val).name),
-    RoborockDyadDataProtocol.CLEAN_MODE: DyadProtocolCacheEntry(lambda val: DyadCleanMode(val).name),
-    RoborockDyadDataProtocol.SUCTION: DyadProtocolCacheEntry(lambda val: DyadSuction(val).name),
-    RoborockDyadDataProtocol.WATER_LEVEL: DyadProtocolCacheEntry(lambda val: DyadWaterLevel(val).name),
-    RoborockDyadDataProtocol.BRUSH_SPEED: DyadProtocolCacheEntry(lambda val: DyadBrushSpeed(val).name),
-    RoborockDyadDataProtocol.POWER: DyadProtocolCacheEntry(lambda val: int(val)),
-    RoborockDyadDataProtocol.AUTO_DRY: DyadProtocolCacheEntry(lambda val: bool(val)),
-    RoborockDyadDataProtocol.MESH_LEFT: DyadProtocolCacheEntry(lambda val: int(360000 - val * 60)),
-    RoborockDyadDataProtocol.BRUSH_LEFT: DyadProtocolCacheEntry(lambda val: int(360000 - val * 60)),
-    RoborockDyadDataProtocol.ERROR: DyadProtocolCacheEntry(lambda val: DyadError(val).name),
-    RoborockDyadDataProtocol.VOLUME_SET: DyadProtocolCacheEntry(lambda val: int(val)),
-    RoborockDyadDataProtocol.STAND_LOCK_AUTO_RUN: DyadProtocolCacheEntry(lambda val: bool(val)),
-    RoborockDyadDataProtocol.AUTO_DRY_MODE: DyadProtocolCacheEntry(lambda val: bool(val)),
-    RoborockDyadDataProtocol.SILENT_DRY_DURATION: DyadProtocolCacheEntry(lambda val: int(val)),  # in minutes
-    RoborockDyadDataProtocol.SILENT_MODE: DyadProtocolCacheEntry(lambda val: bool(val)),
-    RoborockDyadDataProtocol.SILENT_MODE_START_TIME: DyadProtocolCacheEntry(
+    RoborockDyadDataProtocol.STATUS: A01ProtocolCacheEntry(lambda val: RoborockDyadStateCode(val).name),
+    RoborockDyadDataProtocol.SELF_CLEAN_MODE: A01ProtocolCacheEntry(lambda val: DyadSelfCleanMode(val).name),
+    RoborockDyadDataProtocol.SELF_CLEAN_LEVEL: A01ProtocolCacheEntry(lambda val: DyadSelfCleanLevel(val).name),
+    RoborockDyadDataProtocol.WARM_LEVEL: A01ProtocolCacheEntry(lambda val: DyadWarmLevel(val).name),
+    RoborockDyadDataProtocol.CLEAN_MODE: A01ProtocolCacheEntry(lambda val: DyadCleanMode(val).name),
+    RoborockDyadDataProtocol.SUCTION: A01ProtocolCacheEntry(lambda val: DyadSuction(val).name),
+    RoborockDyadDataProtocol.WATER_LEVEL: A01ProtocolCacheEntry(lambda val: DyadWaterLevel(val).name),
+    RoborockDyadDataProtocol.BRUSH_SPEED: A01ProtocolCacheEntry(lambda val: DyadBrushSpeed(val).name),
+    RoborockDyadDataProtocol.POWER: A01ProtocolCacheEntry(lambda val: int(val)),
+    RoborockDyadDataProtocol.AUTO_DRY: A01ProtocolCacheEntry(lambda val: bool(val)),
+    RoborockDyadDataProtocol.MESH_LEFT: A01ProtocolCacheEntry(lambda val: int(360000 - val * 60)),
+    RoborockDyadDataProtocol.BRUSH_LEFT: A01ProtocolCacheEntry(lambda val: int(360000 - val * 60)),
+    RoborockDyadDataProtocol.ERROR: A01ProtocolCacheEntry(lambda val: DyadError(val).name),
+    RoborockDyadDataProtocol.VOLUME_SET: A01ProtocolCacheEntry(lambda val: int(val)),
+    RoborockDyadDataProtocol.STAND_LOCK_AUTO_RUN: A01ProtocolCacheEntry(lambda val: bool(val)),
+    RoborockDyadDataProtocol.AUTO_DRY_MODE: A01ProtocolCacheEntry(lambda val: bool(val)),
+    RoborockDyadDataProtocol.SILENT_DRY_DURATION: A01ProtocolCacheEntry(lambda val: int(val)),  # in minutes
+    RoborockDyadDataProtocol.SILENT_MODE: A01ProtocolCacheEntry(lambda val: bool(val)),
+    RoborockDyadDataProtocol.SILENT_MODE_START_TIME: A01ProtocolCacheEntry(
         lambda val: time(hour=int(val / 60), minute=val % 60)
     ),  # in minutes since 00:00
-    RoborockDyadDataProtocol.SILENT_MODE_END_TIME: DyadProtocolCacheEntry(
+    RoborockDyadDataProtocol.SILENT_MODE_END_TIME: A01ProtocolCacheEntry(
         lambda val: time(hour=int(val / 60), minute=val % 60)
     ),  # in minutes since 00:00
-    RoborockDyadDataProtocol.RECENT_RUN_TIME: DyadProtocolCacheEntry(
+    RoborockDyadDataProtocol.RECENT_RUN_TIME: A01ProtocolCacheEntry(
         lambda val: [int(v) for v in val.split(",")]
     ),  # minutes of cleaning in past few days.
-    RoborockDyadDataProtocol.TOTAL_RUN_TIME: DyadProtocolCacheEntry(lambda val: int(val)),
-    RoborockDyadDataProtocol.SND_STATE: DyadProtocolCacheEntry(lambda val: DyadSndState.from_dict(val)),
-    RoborockDyadDataProtocol.PRODUCT_INFO: DyadProtocolCacheEntry(lambda val: DyadProductInfo.from_dict(val)),
+    RoborockDyadDataProtocol.TOTAL_RUN_TIME: A01ProtocolCacheEntry(lambda val: int(val)),
+    RoborockDyadDataProtocol.SND_STATE: A01ProtocolCacheEntry(lambda val: DyadSndState.from_dict(val)),
+    RoborockDyadDataProtocol.PRODUCT_INFO: A01ProtocolCacheEntry(lambda val: DyadProductInfo.from_dict(val)),
+}
+
+zeo_data_protocol_entries = {
+    # ro
+    RoborockZeoProtocol.STATE: A01ProtocolCacheEntry(lambda val: ZeoState(val).name),
+    RoborockZeoProtocol.COUNTDOWN: A01ProtocolCacheEntry(lambda val: int(val)),
+    RoborockZeoProtocol.WASHING_LEFT: A01ProtocolCacheEntry(lambda val: int(val)),
+    RoborockZeoProtocol.ERROR: A01ProtocolCacheEntry(lambda val: ZeoError(val).name),
+    RoborockZeoProtocol.TIMES_AFTER_CLEAN: A01ProtocolCacheEntry(lambda val: int(val)),
+    RoborockZeoProtocol.DETERGENT_EMPTY: A01ProtocolCacheEntry(lambda val: bool(val)),
+    RoborockZeoProtocol.SOFTENER_EMPTY: A01ProtocolCacheEntry(lambda val: bool(val)),
+    # rw
+    RoborockZeoProtocol.MODE: A01ProtocolCacheEntry(lambda val: ZeoMode(val).name),
+    RoborockZeoProtocol.PROGRAM: A01ProtocolCacheEntry(lambda val: ZeoProgram(val).name),
+    RoborockZeoProtocol.TEMP: A01ProtocolCacheEntry(lambda val: ZeoTemperature(val).name),
+    RoborockZeoProtocol.RINSE_TIMES: A01ProtocolCacheEntry(lambda val: ZeoRinse(val).name),
+    RoborockZeoProtocol.SPIN_LEVEL: A01ProtocolCacheEntry(lambda val: ZeoSpin(val).name),
+    RoborockZeoProtocol.DRYING_MODE: A01ProtocolCacheEntry(lambda val: ZeoDryingMode(val).name),
+    RoborockZeoProtocol.DETERGENT_TYPE: A01ProtocolCacheEntry(lambda val: ZeoDetergentType(val).name),
+    RoborockZeoProtocol.SOFTENER_TYPE: A01ProtocolCacheEntry(lambda val: ZeoSoftenerType(val).name),
+    RoborockZeoProtocol.SOUND_SET: A01ProtocolCacheEntry(lambda val: bool(val)),
 }
 
 
 class RoborockClientA01(RoborockClient):
-    def __init__(self, endpoint: str, device_info: DeviceData):
-        super().__init__(endpoint, device_info)
+    def __init__(self, endpoint: str, device_info: DeviceData, category: RoborockCategory, queue_timeout: int = 4):
+        super().__init__(endpoint, device_info, queue_timeout)
+        self.category = category
 
     def on_message_received(self, messages: list[RoborockMessage]) -> None:
         for message in messages:
@@ -87,14 +120,23 @@ def on_message_received(self, messages: list[RoborockMessage]) -> None:
                     continue
                 payload_json = json.loads(payload.decode())
                 for data_point_number, data_point in payload_json.get("dps").items():
-                    data_point_protocol = RoborockDyadDataProtocol(int(data_point_number))
-                    if data_point_protocol in protocol_entries:
+                    data_point_protocol: RoborockDyadDataProtocol | RoborockZeoProtocol
+                    entries: dict
+                    if self.category == RoborockCategory.WET_DRY_VAC:
+                        data_point_protocol = RoborockDyadDataProtocol(int(data_point_number))
+                        entries = protocol_entries
+                    elif self.category == RoborockCategory.WASHING_MACHINE:
+                        data_point_protocol = RoborockZeoProtocol(int(data_point_number))
+                        entries = zeo_data_protocol_entries
+                    else:
+                        continue
+                    if data_point_protocol in entries:
                         # Auto convert into data struct we want.
-                        converted_response = protocol_entries[data_point_protocol].post_process_fn(data_point)
+                        converted_response = entries[data_point_protocol].post_process_fn(data_point)
                         queue = self._waiting_queue.get(int(data_point_number))
                         if queue and queue.protocol == protocol:
                             queue.resolve((converted_response, None))
 
-    async def update_values(self, dyad_data_protocols: list[RoborockDyadDataProtocol]):
+    async def update_values(self, dyad_data_protocols: list[RoborockDyadDataProtocol | RoborockZeoProtocol]):
         """This should handle updating for each given protocol."""
         raise NotImplementedError
diff --git a/roborock/version_a01_apis/roborock_mqtt_client_a01.py b/roborock/version_a01_apis/roborock_mqtt_client_a01.py
