fix: local_api should receive ip for each device

Author: humbertogontijo

roborock/api.py

@@ -42,6 +42,7 @@
     RoborockDeviceProp,
     RoborockCommand,
     RoborockDockSummary,
+    RoborockDeviceInfo,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -92,8 +93,8 @@ async def request(
 
 class RoborockClient:
 
-    def __init__(self, endpoint: str, device_localkey: dict[str, str]) -> None:
-        self.device_localkey = device_localkey
+    def __init__(self, endpoint: str, devices_info: dict[str, RoborockDeviceInfo]) -> None:
+        self.devices_info = devices_info
         self._seq = 1
         self._random = 4711
         self._id_counter = 10000
@@ -146,7 +147,7 @@ def _decode_msg(self, msg: bytes, local_key: str) -> dict[str, Any]:
         }
 
     def _encode_msg(self, device_id, request_id, protocol, timestamp, payload, prefix=None) -> bytes:
-        local_key = self.device_localkey[device_id]
+        local_key = self.devices_info[device_id].device.local_key
         aes_key = md5bin(encode_timestamp(timestamp) + local_key + self._salt)
         cipher = AES.new(aes_key, AES.MODE_ECB)
         encrypted = cipher.encrypt(pad(payload, AES.block_size))
@@ -172,7 +173,7 @@ def _encode_msg(self, device_id, request_id, protocol, timestamp, payload, prefi
 
     async def on_message(self, device_id, msg) -> bool:
         try:
-            data = self._decode_msg(msg, self.device_localkey[device_id])
+            data = self._decode_msg(msg, self.devices_info[device_id].device.local_key)
             protocol = data.get("protocol")
             if protocol == 102 or protocol == 4:
                 payload = json.loads(data.get("payload").decode())

roborock/containers.py

@@ -229,6 +229,14 @@ class WashTowelField(str, Enum):
     Wash_MODE = "wash_mode"
 
 
+class NetworkInfoField(str, Enum):
+    SSID = "ssid"
+    IP = "ip"
+    MAC = "mac"
+    BSSID = "bssid"
+    RSSI = "rssi"
+
+
 class RoborockBase(dict):
     def __init__(self, data: dict[str, any]) -> None:
         super().__init__()
@@ -1044,3 +1052,23 @@ def smart_wash(self) -> int:
     @property
     def wash_interval(self) -> int:
         return self.get(SmartWashField.WASH_INTERVAL)
+
+class NetworkInfo(RoborockBase):
+    def __init__(self, data: dict[str, any]) -> None:
+        super().__init__(data)
+
+    @property
+    def ssid(self) -> str:
+        return self.get(NetworkInfoField.SSID)
+    @property
+    def ip(self) -> str:
+        return self.get(NetworkInfoField.IP)
+    @property
+    def mac(self) -> str:
+        return self.get(NetworkInfoField.MAC)
+    @property
+    def bssid(self) -> str:
+        return self.get(NetworkInfoField.BSSID)
+    @property
+    def rssi(self) -> int:
+        return self.get(NetworkInfoField.RSSI)

roborock/local_api.py

@@ -3,14 +3,14 @@
 import asyncio
 import logging
 import socket
-from asyncio import Lock
+from asyncio import Lock, Future
 from typing import Callable, Coroutine, Any
 
 import async_timeout
 
 from roborock.api import RoborockClient, SPECIAL_COMMANDS
 from roborock.exceptions import RoborockTimeout, CommandVacuumError
-from roborock.typing import RoborockCommand
+from roborock.typing import RoborockCommand, RoborockDeviceInfo
 from roborock.util import get_running_loop_or_create_one
 
 secured_prefix = 199
@@ -21,13 +21,26 @@
 _LOGGER = logging.getLogger(__name__)
 
 
+class RoborockProtocol(asyncio.DatagramProtocol):
+
+    def __init__(self, fut: Future):
+        self.fut = fut
+
+    def datagram_received(self, data, addr):
+        self.fut.set_result((data, addr))
+
+    def error_received(self, exc):
+        self.fut.set_exception(exc)
+
+
 class RoborockLocalClient(RoborockClient):
 
-    def __init__(self, ip: str, device_localkey: dict[str, str]):
-        super().__init__("abc", device_localkey)
+    def __init__(self, devices_info: dict[str, RoborockDeviceInfo]):
+        super().__init__("abc", devices_info)
+        self.loop = get_running_loop_or_create_one()
         self.device_listener: dict[str, RoborockSocketListener] = {
-            device_id: RoborockSocketListener(ip, device_id, self.on_message)
-            for device_id in device_localkey
+            device_id: RoborockSocketListener(device_info.network_info.ip, device_id, self.on_message)
+            for device_id, device_info in devices_info.items()
         }
         self._mutex = Lock()
 

roborock/offline/offline.py

@@ -1,20 +1,30 @@
 import asyncio
 import logging
 
+from roborock.containers import HomeDataDevice, HomeDataDeviceField, HomeDataProduct, HomeDataProductField, NetworkInfo, \
+    NetworkInfoField
 from roborock.local_api import RoborockLocalClient
+from roborock.typing import RoborockDeviceInfo
 
-local_ip = "<local_ip>"
-local_key = "<local_key>"
+local_ip = "192.168.1.232"
+local_key = "nXTBj42ej5WxQopO"
+device_id = "1r9W0cAmDZ2COuVekgRhKA"
 
 
 async def main():
     logging_config = {
         "level": logging.DEBUG
     }
     logging.basicConfig(**logging_config)
-    device_id = "1r9W0cAmDZ2COuVekgRhKA"
-    client = RoborockLocalClient(local_ip, {
-        "1r9W0cAmDZ2COuVekgRhKA": local_key
+    client = RoborockLocalClient({
+        device_id: RoborockDeviceInfo(HomeDataDevice({
+            HomeDataDeviceField.DUID: device_id,
+            HomeDataDeviceField.LOCAL_KEY: local_key
+        }), HomeDataProduct({
+            HomeDataProductField.MODEL: "test"
+        }), NetworkInfo({
+            NetworkInfoField.IP: local_ip
+        }))
     })
     await client.async_connect()
     props = await client.get_prop(device_id)

roborock/typing.py

@@ -1,8 +1,10 @@
+from __future__ import annotations
+
 from enum import Enum
 
 from .code_mappings import RoborockDockDustCollectionType, RoborockDockWashingModeType
 from .containers import Status, CleanSummary, Consumable, \
-    DNDTimer, CleanRecord, SmartWashParameters
+    DNDTimer, CleanRecord, SmartWashParameters, HomeDataDevice, HomeDataProduct, NetworkInfo
 
 
 class RoborockDevicePropField(str, Enum):
@@ -84,7 +86,13 @@ class RoborockCommand(str, Enum):
     SET_SERVER_TIMER = "set_server_timer"
     APP_GET_INIT_STATUS = "get_init_status"
     SET_APP_TIMEZONE = "set_app_timezone"
+    GET_NETWORK_INFO = "get_network_info"
 
+class RoborockDeviceInfo:
+    def __init__(self, device: HomeDataDevice, product: HomeDataProduct, network_info: NetworkInfo = None):
+        self.device = device
+        self.product = product
+        self.network_info: NetworkInfo | None = network_info
 
 class RoborockDockSummary:
     def __init__(self, dust_collection_mode: RoborockDockDustCollectionType,