python-roborock 2.44.1__tar.gz → 2.46.0__tar.gz

{python_roborock-2.44.1 → python_roborock-2.46.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: python-roborock
-Version: 2.44.1
+Version: 2.46.0
 Summary: A package to control Roborock vacuums.
 Home-page: https://github.com/humbertogontijo/python-roborock
 License: GPL-3.0-only

{python_roborock-2.44.1 → python_roborock-2.46.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "python-roborock"
-version = "2.44.1"
+version = "2.46.0"
 description = "A package to control Roborock vacuums."
 authors = ["humbertogontijo <humbertogontijo@users.noreply.github.com>"]
 license = "GPL-3.0-only"

{python_roborock-2.44.1 → python_roborock-2.46.0}/roborock/api.py

@@ -23,7 +23,7 @@ from .roborock_message import (
 from .util import get_next_int
 
 _LOGGER = logging.getLogger(__name__)
-KEEPALIVE = 60
+KEEPALIVE = 70
 
 
 class RoborockClient(ABC):
@@ -78,12 +78,6 @@ class RoborockClient(ABC):
             return False
         return True
 
-    async def validate_connection(self) -> None:
-        if not self.should_keepalive():
-            self._logger.info("Resetting Roborock connection due to keepalive timeout")
-            await self.async_disconnect()
-        await self.async_connect()
-
     async def _wait_response(self, request_id: int, queue: RoborockFuture) -> Any:
         try:
             response = await queue.async_get(self.queue_timeout)

{python_roborock-2.44.1 → python_roborock-2.46.0}/roborock/cloud_api.py

@@ -8,6 +8,10 @@ from asyncio import Lock
 from typing import Any
 
 import paho.mqtt.client as mqtt
+from paho.mqtt.enums import MQTTErrorCode
+
+# Mypy is not seeing this for some reason. It wants me to use the depreciated ReasonCodes
+from paho.mqtt.reasoncodes import ReasonCode  # type: ignore
 
 from .api import KEEPALIVE, RoborockClient
 from .containers import DeviceData, UserData
@@ -67,7 +71,8 @@ class RoborockMqttClient(RoborockClient, ABC):
         self._mqtt_client = _Mqtt()
         self._mqtt_client.on_connect = self._mqtt_on_connect
         self._mqtt_client.on_message = self._mqtt_on_message
-        self._mqtt_client.on_disconnect = self._mqtt_on_disconnect
+        # Due to the incorrect ReasonCode, it is confused by typing
+        self._mqtt_client.on_disconnect = self._mqtt_on_disconnect  # type: ignore
         if mqtt_params.tls:
             self._mqtt_client.tls_set()
 
@@ -76,12 +81,20 @@ class RoborockMqttClient(RoborockClient, ABC):
         self._mutex = Lock()
         self._decoder: Decoder = create_mqtt_decoder(device_info.device.local_key)
         self._encoder: Encoder = create_mqtt_encoder(device_info.device.local_key)
+        self.received_message_since_last_disconnect = False
+        self._topic = f"rr/m/o/{self._mqtt_user}/{self._hashed_user}/{self.device_info.device.duid}"
 
-    def _mqtt_on_connect(self, *args, **kwargs):
-        _, __, ___, rc, ____ = args
+    def _mqtt_on_connect(
+        self,
+        client: mqtt.Client,
+        userdata: object,
+        flags: dict[str, int],
+        rc: ReasonCode,
+        properties: mqtt.Properties | None = None,
+    ):
         connection_queue = self._waiting_queue.get(CONNECT_REQUEST_ID)
-        if rc != mqtt.MQTT_ERR_SUCCESS:
-            message = f"Failed to connect ({mqtt.error_string(rc)})"
+        if rc.is_failure:
+            message = f"Failed to connect ({rc})"
             self._logger.error(message)
             if connection_queue:
                 connection_queue.set_exception(VacuumError(message))
@@ -89,19 +102,19 @@ class RoborockMqttClient(RoborockClient, ABC):
                 self._logger.debug("Failed to notify connect future, not in queue")
             return
         self._logger.info(f"Connected to mqtt {self._mqtt_host}:{self._mqtt_port}")
-        topic = f"rr/m/o/{self._mqtt_user}/{self._hashed_user}/{self.device_info.device.duid}"
-        (result, mid) = self._mqtt_client.subscribe(topic)
+        (result, mid) = self._mqtt_client.subscribe(self._topic)
         if result != 0:
-            message = f"Failed to subscribe ({mqtt.error_string(rc)})"
+            message = f"Failed to subscribe ({str(rc)})"
             self._logger.error(message)
             if connection_queue:
                 connection_queue.set_exception(VacuumError(message))
             return
-        self._logger.info(f"Subscribed to topic {topic}")
+        self._logger.info(f"Subscribed to topic {self._topic}")
         if connection_queue:
             connection_queue.set_result(True)
 
     def _mqtt_on_message(self, *args, **kwargs):
+        self.received_message_since_last_disconnect = True
         client, __, msg = args
         try:
             messages = self._decoder(msg.payload)
@@ -109,10 +122,16 @@ class RoborockMqttClient(RoborockClient, ABC):
         except Exception as ex:
             self._logger.exception(ex)
 
-    def _mqtt_on_disconnect(self, *args, **kwargs):
-        _, __, rc, ___ = args
+    def _mqtt_on_disconnect(
+        self,
+        client: mqtt.Client,
+        data: object,
+        flags: dict[str, int],
+        rc: ReasonCode | None,
+        properties: mqtt.Properties | None = None,
+    ):
         try:
-            exc = RoborockException(mqtt.error_string(rc)) if rc != mqtt.MQTT_ERR_SUCCESS else None
+            exc = RoborockException(str(rc)) if rc is not None and rc.is_failure else None
             super().on_connection_lost(exc)
             connection_queue = self._waiting_queue.get(DISCONNECT_REQUEST_ID)
             if connection_queue:
@@ -138,7 +157,7 @@ class RoborockMqttClient(RoborockClient, ABC):
 
         if rc != mqtt.MQTT_ERR_SUCCESS:
             disconnected_future.cancel()
-            raise RoborockException(f"Failed to disconnect ({mqtt.error_string(rc)})")
+            raise RoborockException(f"Failed to disconnect ({str(rc)})")
 
         return disconnected_future
 
@@ -178,3 +197,63 @@ class RoborockMqttClient(RoborockClient, ABC):
         )
         if info.rc != mqtt.MQTT_ERR_SUCCESS:
             raise RoborockException(f"Failed to publish ({mqtt.error_string(info.rc)})")
+
+    async def _unsubscribe(self) -> MQTTErrorCode:
+        """Unsubscribe from the topic."""
+        loop = asyncio.get_running_loop()
+        (result, mid) = await loop.run_in_executor(None, self._mqtt_client.unsubscribe, self._topic)
+
+        if result != 0:
+            message = f"Failed to unsubscribe ({mqtt.error_string(result)})"
+            self._logger.error(message)
+        else:
+            self._logger.info(f"Unsubscribed from topic {self._topic}")
+        return result
+
+    async def _subscribe(self) -> MQTTErrorCode:
+        """Subscribe to the topic."""
+        loop = asyncio.get_running_loop()
+        (result, mid) = await loop.run_in_executor(None, self._mqtt_client.subscribe, self._topic)
+
+        if result != 0:
+            message = f"Failed to subscribe ({mqtt.error_string(result)})"
+            self._logger.error(message)
+        else:
+            self._logger.info(f"Subscribed to topic {self._topic}")
+        return result
+
+    async def _reconnect(self) -> None:
+        """Reconnect to the MQTT broker."""
+        await self.async_disconnect()
+        await self.async_connect()
+
+    async def _validate_connection(self) -> None:
+        """Override the default validate connection to try to re-subscribe rather than disconnect.
+        When something seems to be wrong with our connection, we should follow the following steps:
+        1. Try to unsubscribe and resubscribe from the topic.
+        2. If we don't end up getting a message, we should completely disconnect and reconnect to the MQTT broker.
+        3. We will continue to try to disconnect and reconnect until we get a message.
+        4. If we get a message, the next time connection is lost, We will go back to step 1.
+        """
+        # If we should no longer keep the current connection alive...
+        if not self.should_keepalive():
+            self._logger.info("Resetting Roborock connection due to keepalive timeout")
+            if not self.received_message_since_last_disconnect:
+                # If we have already tried to unsub and resub, and we are still in this state,
+                # we should try to reconnect.
+                return await self._reconnect()
+            try:
+                # Mark that we have tried to unsubscribe and resubscribe
+                self.received_message_since_last_disconnect = False
+                if await self._unsubscribe() != 0:
+                    # If we fail to unsubscribe, reconnect to the broker
+                    return await self._reconnect()
+                if await self._subscribe() != 0:
+                    # If we fail to subscribe, reconnected to the broker.
+                    return await self._reconnect()
+
+            except Exception:  # noqa
+                # If we get any errors at all, we should just reconnect.
+                return await self._reconnect()
+        # Call connect to make sure everything is still in a good state.
+        await self.async_connect()

{python_roborock-2.44.1 → python_roborock-2.46.0}/roborock/exceptions.py

@@ -1,4 +1,5 @@
 """Roborock exceptions."""
+
 from __future__ import annotations
 
 

{python_roborock-2.44.1 → python_roborock-2.46.0}/roborock/version_1_apis/roborock_local_client_v1.py

@@ -138,6 +138,12 @@ class RoborockLocalClientV1(RoborockClientV1, RoborockClient):
             response_protocol=RoborockMessageProtocol.PING_RESPONSE,
         )
 
+    async def _validate_connection(self) -> None:
+        if not self.should_keepalive():
+            self._logger.info("Resetting Roborock connection due to keepalive timeout")
+            await self.async_disconnect()
+        await self.async_connect()
+
     def _send_msg_raw(self, data: bytes):
         try:
             if not self.transport:
@@ -172,7 +178,7 @@ class RoborockLocalClientV1(RoborockClientV1, RoborockClient):
         method: str | None = None,
         params: list | dict | int | None = None,
     ) -> RoborockMessage:
-        await self.validate_connection()
+        await self._validate_connection()
         msg = self._encoder(roborock_message)
         if method:
             self._logger.debug(f"id={request_id} Requesting method {method} with {params}")

{python_roborock-2.44.1 → python_roborock-2.46.0}/roborock/version_1_apis/roborock_mqtt_client_v1.py

@@ -51,7 +51,7 @@ class RoborockMqttClientV1(RoborockMqttClient, RoborockClientV1):
         )
         self._logger.debug("Building message id %s for method %s", request_message.request_id, method)
 
-        await self.validate_connection()
+        await self._validate_connection()
         request_id = request_message.request_id
         response_protocol = (
             RoborockMessageProtocol.MAP_RESPONSE if method in COMMANDS_SECURED else RoborockMessageProtocol.RPC_RESPONSE

{python_roborock-2.44.1 → python_roborock-2.46.0}/roborock/version_a01_apis/roborock_mqtt_client_a01.py

@@ -40,7 +40,7 @@ class RoborockMqttClientA01(RoborockMqttClient, RoborockClientA01):
         self._logger = RoborockLoggerAdapter(device_info.device.name, _LOGGER)
 
     async def _send_message(self, roborock_message: RoborockMessage):
-        await self.validate_connection()
+        await self._validate_connection()
         response_protocol = RoborockMessageProtocol.RPC_RESPONSE
 
         m = self._encoder(roborock_message)

{python_roborock-2.44.1 → python_roborock-2.46.0}/roborock/web_api.py

@@ -6,6 +6,7 @@ import hmac
 import logging
 import math
 import secrets
+import string
 import time
 
 import aiohttp
@@ -190,6 +191,113 @@ class RoborockApiClient:
             else:
                 raise RoborockException(f"{code_response.get('msg')} - response code: {code_response.get('code')}")
 
+    async def request_code_v4(self) -> None:
+        """Request a code using the v4 endpoint."""
+        try:
+            self._login_limiter.try_acquire("login")
+        except BucketFullException as ex:
+            _LOGGER.info(ex.meta_info)
+            raise RoborockRateLimit("Reached maximum requests for login. Please try again later.") from ex
+        base_url = await self._get_base_url()
+        header_clientid = self._get_header_client_id()
+        code_request = PreparedRequest(
+            base_url,
+            self.session,
+            {
+                "header_clientid": header_clientid,
+                "Content-Type": "application/x-www-form-urlencoded",
+                "header_clientlang": "en",
+            },
+        )
+
+        code_response = await code_request.request(
+            "post",
+            "/api/v4/email/code/send",
+            params={"email": self._username, "type": "login", "platform": ""},
+        )
+        if code_response is None:
+            raise RoborockException("Failed to get a response from send email code")
+        response_code = code_response.get("code")
+        if response_code != 200:
+            _LOGGER.info("Request code failed for %s with the following context: %s", self._username, code_response)
+            if response_code == 2008:
+                raise RoborockAccountDoesNotExist("Account does not exist - check your login and try again.")
+            elif response_code == 9002:
+                raise RoborockTooFrequentCodeRequests("You have attempted to request too many codes. Try again later")
+            else:
+                raise RoborockException(f"{code_response.get('msg')} - response code: {code_response.get('code')}")
+
+    async def sign_key_v3(self, s: str) -> str:
+        """Sign a randomly generated string."""
+        base_url = await self._get_base_url()
+        header_clientid = self._get_header_client_id()
+        code_request = PreparedRequest(base_url, self.session, {"header_clientid": header_clientid})
+
+        code_response = await code_request.request(
+            "post",
+            "/api/v3/key/sign",
+            params={"s": s},
+        )
+
+        if not code_response or "data" not in code_response or "k" not in code_response["data"]:
+            raise RoborockException("Failed to get a response from sign key")
+        response_code = code_response.get("code")
+
+        if response_code != 200:
+            _LOGGER.info("Request code failed for %s with the following context: %s", self._username, code_response)
+            raise RoborockException(f"{code_response.get('msg')} - response code: {code_response.get('code')}")
+
+        return code_response["data"]["k"]
+
+    async def code_login_v4(self, code: int | str, country: str, country_code: int) -> UserData:
+        """
+        Login via code authentication.
+        :param code: The code from the email.
+        :param country: The two-character representation of the country, i.e. "US"
+        :param country_code: the country phone number code i.e. 1 for US.
+        """
+        base_url = await self._get_base_url()
+        header_clientid = self._get_header_client_id()
+        x_mercy_ks = "".join(secrets.choice(string.ascii_letters + string.digits) for _ in range(16))
+        x_mercy_k = await self.sign_key_v3(x_mercy_ks)
+        login_request = PreparedRequest(
+            base_url,
+            self.session,
+            {"header_clientid": header_clientid, "x-mercy-ks": x_mercy_ks, "x-mercy-k": x_mercy_k},
+        )
+        login_response = await login_request.request(
+            "post",
+            "/api/v4/auth/email/login/code",
+            params={
+                "country": country,
+                "countryCode": country_code,
+                "email": self._username,
+                "code": code,
+                # Major and minor version are the user agreement version, we will need to see if this needs to be
+                # dynamic https://usiot.roborock.com/api/v3/app/agreement/latest?country=US
+                "majorVersion": 14,
+                "minorVersion": 0,
+            },
+        )
+        if login_response is None:
+            raise RoborockException("Login request response is None")
+        response_code = login_response.get("code")
+        if response_code != 200:
+            _LOGGER.info("Login failed for %s with the following context: %s", self._username, login_response)
+            if response_code == 2018:
+                raise RoborockInvalidCode("Invalid code - check your code and try again.")
+            if response_code == 3009:
+                raise RoborockNoUserAgreement("You must accept the user agreement in the Roborock app to continue.")
+            if response_code == 3006:
+                raise RoborockInvalidUserAgreement(
+                    "User agreement must be accepted again - or you are attempting to use the Mi Home app account."
+                )
+            raise RoborockException(f"{login_response.get('msg')} - response code: {response_code}")
+        user_data = login_response.get("data")
+        if not isinstance(user_data, dict):
+            raise RoborockException("Got unexpected data type for user_data")
+        return UserData.from_dict(user_data)
+
     async def pass_login(self, password: str) -> UserData:
         try:
             self._login_limiter.try_acquire("login")