chore: random changes

Author: Lash-L

roborock/mqtt_manager.py

@@ -4,28 +4,39 @@
 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
 
 _LOGGER = logging.getLogger(__name__)
 
 
 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:

roborock/roborock_device.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)