fix: minor fixes

humbertogontijo · humbertogontijo · commit 29bdb4542e1c · 2023-03-31T22:12:47.000-03:00
diff --git a/roborock/api.py b/roborock/api.py
@@ -47,7 +47,9 @@
 _LOGGER = logging.getLogger(__name__)
 QUEUE_TIMEOUT = 4
 MQTT_KEEPALIVE = 60
-
+SPECIAL_COMMANDS = [
+    RoborockCommand.GET_MAP_V1,
+]
 
 def md5hex(message: str) -> str:
     md5 = hashlib.md5()
@@ -119,6 +121,12 @@ def _decode_msg(self, msg: bytes, local_key: str) -> dict[str, Any]:
         [version, _seq, _random, timestamp, protocol, payload_len] = struct.unpack(
             "!3sIIIHH", msg[0:19]
         )
+        if not payload_len:
+            return {
+                "version": version,
+                "timestamp": timestamp,
+                "protocol": protocol,
+            }
         [payload, expected_crc32] = struct.unpack_from(f"!{payload_len}sI", msg, 19)
         # if crc32 != expected_crc32:
         #     raise RoborockException(f"Wrong CRC32 {crc32}, expected {expected_crc32}")
@@ -158,7 +166,7 @@ def _encode_msg(self, device_id, protocol, timestamp, payload, prefix='') -> byt
         msg += struct.pack("!I", crc32)
         return msg
 
-    async def on_message(self, device_id, msg) -> None:
+    async def on_message(self, device_id, msg) -> bool:
         try:
             data = self._decode_msg(msg, self.device_localkey[device_id])
             protocol = data.get("protocol")
@@ -184,11 +192,12 @@ async def on_message(self, device_id, msg) -> None:
                                     )
                                 else:
                                     result = data_point_response.get("result")
-                                    if isinstance(result, list) and len(result) > 0:
+                                    if isinstance(result, list) and len(result) == 1:
                                         result = result[0]
                                     await queue.async_put(
                                         (result, None), timeout=QUEUE_TIMEOUT
                                     )
+                                    return True
                         elif request_id < self._id_counter:
                             _LOGGER.debug(
                                 f"id={request_id} Ignoring response: {data_point_response}"
@@ -217,6 +226,7 @@ async def on_message(self, device_id, msg) -> None:
                         if isinstance(decrypted, list):
                             decrypted = decrypted[0]
                         await queue.async_put((decrypted, None), timeout=QUEUE_TIMEOUT)
+                        return True
         except Exception as ex:
             _LOGGER.exception(ex)
 
@@ -234,7 +244,7 @@ async def _async_response(self, request_id: int, protocol_id: int = 0) -> tuple[
             del self._waiting_queue[request_id]
 
     def _get_payload(
-            self, method: RoborockCommand, params: list = None
+            self, method: RoborockCommand, params: list = None, secured=False
     ):
         timestamp = math.floor(time.time())
         request_id = self._id_counter
@@ -243,11 +253,12 @@ def _get_payload(
             "id": request_id,
             "method": method,
             "params": params or [],
-            "security": {
-                "endpoint": self._endpoint,
-                "nonce": self._nonce.hex().upper(),
-            },
         }
+        if secured:
+            inner["security"] = {
+                "endpoint": self._endpoint,
+                "nonce": "39344139463753454b4851444f4a4442",
+            }
         payload = bytes(
             json.dumps(
                 {
@@ -328,10 +339,10 @@ async def get_prop(self, device_id: str) -> RoborockDeviceProp:
             ]
         )
         last_clean_record = None
-        if clean_summary and clean_summary.records and len(clean_summary.records) > 0:
-            last_clean_record = await self.get_clean_record(
-                device_id, clean_summary.records[0]
-            )
+        # if clean_summary and clean_summary.records and len(clean_summary.records) > 0:
+        #     last_clean_record = await self.get_clean_record(
+        #         device_id, clean_summary.records[0]
+        #     )
         dock_summary = None
         if status and status.dock_type != RoborockDockType.NO_DOCK:
             dock_summary = await self.get_dock_summary(device_id, status.dock_type)
diff --git a/roborock/cloud_api.py b/roborock/cloud_api.py
@@ -10,7 +10,7 @@
 
 import paho.mqtt.client as mqtt
 
-from roborock.api import md5hex, md5bin, RoborockClient
+from roborock.api import md5hex, md5bin, RoborockClient, SPECIAL_COMMANDS
 from roborock.exceptions import (
     RoborockException,
     CommandVacuumError,
@@ -28,9 +28,6 @@
 _LOGGER = logging.getLogger(__name__)
 QUEUE_TIMEOUT = 4
 MQTT_KEEPALIVE = 60
-COMMANDS_WITH_BINARY_RESPONSE = [
-    RoborockCommand.GET_MAP_V1,
-]
 
 
 class RoborockMqttClient(RoborockClient, mqtt.Client):
@@ -192,10 +189,10 @@ async def send_command(
             self, device_id: str, method: RoborockCommand, params: list = None
     ):
         await self.validate_connection()
-        request_id, timestamp, payload = super()._get_payload(method, params)
+        request_id, timestamp, payload = super()._get_payload(method, params, True)
         _LOGGER.debug(f"id={request_id} Requesting method {method} with {params}")
         request_protocol = 101
-        response_protocol = 301 if method in COMMANDS_WITH_BINARY_RESPONSE else 102
+        response_protocol = 301 if method in SPECIAL_COMMANDS else 102
         msg = super()._encode_msg(device_id, request_protocol, timestamp, payload)
         self._send_msg_raw(device_id, msg)
         (response, err) = await self._async_response(request_id, response_protocol)
diff --git a/roborock/local_api.py b/roborock/local_api.py
@@ -3,16 +3,21 @@
 import asyncio
 import logging
 import socket
+from asyncio import Lock
 from typing import Callable, Coroutine
 
 import async_timeout
 
-from roborock.api import RoborockClient
+from roborock.api import RoborockClient, SPECIAL_COMMANDS
 from roborock.exceptions import RoborockTimeout, CommandVacuumError
 from roborock.typing import RoborockCommand
 from roborock.util import get_running_loop_or_create_one
 
 secured_prefix = 199
+get_prefix = 119
+app_prefix = 135
+set_prefix = 151
+
 _LOGGER = logging.getLogger(__name__)
 
 
@@ -34,9 +39,10 @@ async def async_connect(self):
     async def send_command(
             self, device_id: str, method: RoborockCommand, params: list = None
     ):
-        request_id, timestamp, payload = super()._get_payload(method, params)
+        secured = True if method in SPECIAL_COMMANDS else False
+        request_id, timestamp, payload = self._get_payload(method, params, secured)
         _LOGGER.debug(f"id={request_id} Requesting method {method} with {params}")
-        prefix = secured_prefix
+        prefix = secured_prefix if method in SPECIAL_COMMANDS else get_prefix
         protocol = 4
         msg = self._encode_msg(device_id, protocol, timestamp, payload, prefix)
         _LOGGER.debug(f"Requesting with prefix {prefix} and payload {payload}")
@@ -49,26 +55,38 @@ async def send_command(
         _LOGGER.debug(f"id={request_id} Response from {method}: {response}")
         return response
 
+class RoborockSocket(socket.socket):
+    _closed = None
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @property
+    def is_closed(self):
+        return self._closed
 
 class RoborockSocketListener:
     roborock_port = 58867
 
-    def __init__(self, ip: str, device_id: str, on_message: Callable[[str, bytes], Coroutine | None],
+    def __init__(self, ip: str, device_id: str, on_message: Callable[[str, bytes], Coroutine[bool] | bool],
                  timeout: float | int = 4):
         self.ip = ip
         self.device_id = device_id
-        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.socket = RoborockSocket(socket.AF_INET, socket.SOCK_STREAM)
         self.socket.setblocking(False)
         self.loop = get_running_loop_or_create_one()
         self.on_message = on_message
         self.timeout = timeout
         self.is_connected = False
+        self._lock = Lock()
 
     async def _main_coro(self):
-        while self.is_connected:
+        while not self.socket.is_closed:
             try:
                 message = await self.loop.sock_recv(self.socket, 4096)
-                await self.on_message(self.device_id, message)
+                accepted = await self.on_message(self.device_id, message)
+                if accepted:
+                    self._lock.release() if self._lock.locked() else None
             except Exception as e:
                 _LOGGER.exception(e)
                 self.is_connected = False
@@ -78,14 +96,16 @@ async def connect(self):
         async with async_timeout.timeout(self.timeout):
             await self.loop.sock_connect(self.socket, (self.ip, 58867))
             self.is_connected = True
-        asyncio.create_task(self._main_coro())
+        self.loop.create_task(self._main_coro())
 
     async def send_message(self, data: bytes, local_key: str):
         response = {}
         try:
             async with async_timeout.timeout(self.timeout):
+                await self._lock.acquire()
                 await self.loop.sock_sendall(self.socket, data)
         except (asyncio.TimeoutError, asyncio.CancelledError):
+            self._lock.release() if self._lock.locked() else None
             raise RoborockTimeout(
                 f"Timeout after {self.timeout} seconds waiting for response"
             ) from None
