v1.0.2

README.md

@@ -1,7 +1,7 @@
 [![version](https://img.shields.io/github/manifest-json/v/Tasshack/dreame-vacuum?filename=custom_components%2Fdreame_vacuum%2Fmanifest.json&color=green)](https://github.com/Tasshack/dreame-vacuum/releases/latest)
 [![HACS](https://img.shields.io/badge/HACS-Default-orange.svg?logo=HomeAssistantCommunityStore&logoColor=white)](https://github.com/hacs/integration)
 [![Community Forum](https://img.shields.io/static/v1.svg?label=Community&message=Forum&color=41bdf5&logo=HomeAssistant&logoColor=white)](https://community.home-assistant.io/t/custom-component-dreame-vacuum/473026)
-[![Ko-Fi](https://img.shields.io/static/v1.svg?label=%20&message=Ko-Fi&color=F16061&logo=ko-fi&logoColor=white)](https://www.ko-fi/Tasshack)
+[![Ko-Fi](https://img.shields.io/static/v1.svg?label=%20&message=Ko-Fi&color=F16061&logo=ko-fi&logoColor=white)](https://www.ko-fi.com/Tasshack)
 [![PayPal.Me](https://img.shields.io/static/v1.svg?label=%20&message=PayPal.Me&logo=paypal)](https://paypal.me/Tasshackk)
 
 ![dreame Logo](https://cdn.shopify.com/s/files/1/0302/5276/1220/files/rsz_logo_-01_400x_2ecfe8c0-2756-4bd1-a3f4-593b1f73e335_284x.jpg "dreame Logo")

custom_components/dreame_vacuum/config_flow.py

@@ -52,7 +52,7 @@
     "dreame.vacuum.p2187",
     "dreame.vacuum.r2328",
     "dreame.vacuum.p2028a",
-    "dreame.vacuum.r2251a",
+    #"dreame.vacuum.r2251a", Map private key missing
     "dreame.vacuum.p2029",
     "dreame.vacuum.r2257o",
     "dreame.vacuum.r2215o",
@@ -91,7 +91,7 @@
     "dreame.vacuum.p2140a",
     "dreame.vacuum.p2114a",
     "dreame.vacuum.p2114o",
-    "dreame.vacuum.r2210",
+    #"dreame.vacuum.r2210", Map private key missing
     "dreame.vacuum.p2149o",
     "dreame.vacuum.p2150a",
     "dreame.vacuum.p2150b",

custom_components/dreame_vacuum/const.py

@@ -110,6 +110,8 @@
 NOTIFICATION_ID_REPLACE_DETERGENT: Final = "replace_detergent"
 NOTIFICATION_ID_CLEANUP_COMPLETED: Final = "cleanup_completed"
 NOTIFICATION_ID_WARNING: Final = "warning"
+NOTIFICATION_ID_INFORMATION: Final = "information"
+NOTIFICATION_ID_CONSUMABLE: Final = "consumable"
 NOTIFICATION_ID_ERROR: Final = "error"
 NOTIFICATION_ID_REPLACE_TEMPORARY_MAP: Final = "replace_temporary_map"
 NOTIFICATION_ID_2FA_LOGIN: Final = "2fa_login"

custom_components/dreame_vacuum/coordinator.py

@@ -14,6 +14,7 @@
     ATTR_ENTITY_ID
 )
 from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity import generate_entity_id
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 from .dreame import DreameVacuumDevice, DreameVacuumProperty
@@ -52,6 +53,8 @@
     NOTIFICATION_ID_CLEANUP_COMPLETED,
     NOTIFICATION_ID_WARNING,
     NOTIFICATION_ID_ERROR,
+    NOTIFICATION_ID_INFORMATION,
+    NOTIFICATION_ID_CONSUMABLE,
     NOTIFICATION_ID_REPLACE_TEMPORARY_MAP,
     NOTIFICATION_ID_2FA_LOGIN,
     EVENT_TASK_STATUS,
@@ -120,9 +123,9 @@ def __init__(
             LOGGER,
             name=DOMAIN,
         )
-
-        hass.bus.async_listen(
-            persistent_notification.EVENT_PERSISTENT_NOTIFICATIONS_UPDATED,
+        async_dispatcher_connect(
+            hass,
+            persistent_notification.SIGNAL_PERSISTENT_NOTIFICATIONS_UPDATED,
             self._notification_dismiss_listener,
         )
 
@@ -255,52 +258,60 @@ def _has_temporary_map_changed(self, previous_value=None) -> None:
                 NOTIFICATION_ID_REPLACE_TEMPORARY_MAP)
 
     def _create_persistent_notification(self, content, notification_id) -> None:
-        if self._notify:
-            if isinstance(self._notify, list):
+        if self._notify or notification_id == NOTIFICATION_ID_2FA_LOGIN:
+            if isinstance(self._notify, list) and notification_id != NOTIFICATION_ID_2FA_LOGIN:
                 if notification_id == NOTIFICATION_ID_CLEANUP_COMPLETED:
                     if NOTIFICATION_ID_CLEANUP_COMPLETED not in self._notify:
                         return
-                elif NOTIFICATION_ID_WARNING in notification_id:
+                elif (
+                    NOTIFICATION_ID_WARNING in notification_id
+                ):
                     if NOTIFICATION_ID_WARNING not in self._notify:
                         return
                 elif NOTIFICATION_ID_ERROR in notification_id:
                     if NOTIFICATION_ID_ERROR not in self._notify:
                         return
-                elif notification_id == NOTIFICATION_ID_DUST_COLLECTION or notification_id == NOTIFICATION_ID_CLEANING_PAUSED:
-                    if "information" not in self._notify:
-                        return
-                elif notification_id != NOTIFICATION_ID_REPLACE_TEMPORARY_MAP: 
-                    if "consumable" not in self._notify:
+                elif (
+                    notification_id == NOTIFICATION_ID_DUST_COLLECTION
+                    or notification_id == NOTIFICATION_ID_CLEANING_PAUSED
+                ):
+                    if NOTIFICATION_ID_INFORMATION not in self._notify:
                         return
+                elif (
+                    notification_id != NOTIFICATION_ID_REPLACE_TEMPORARY_MAP
+                ):
+                    if NOTIFICATION_ID_CONSUMABLE not in self._notify:
+                       return
+
 
             persistent_notification.async_create(
                 self.hass,
                 content,
                 title=self.device.name,
-                notification_id=f"{DOMAIN}_{notification_id}",
+                notification_id=f"{DOMAIN}_{self.device.mac}_{notification_id}"
             )
 
     def _remove_persistent_notification(self, notification_id) -> None:
         persistent_notification.async_dismiss(
-            self.hass, f"{DOMAIN}_{notification_id}")
-
-    def _notification_dismiss_listener(self, event) -> None:
-        notifications = self.hass.data.get(persistent_notification.DOMAIN)
-
-        if self._has_warning:
-            if (
-                f"{persistent_notification.DOMAIN}.{DOMAIN}_{NOTIFICATION_ID_WARNING}"
-                not in notifications
-            ):
-                self._has_warning = False
-                self.device.clear_warning()
-
-        if self._two_factor_url:
-            if (
-                f"{persistent_notification.DOMAIN}.{DOMAIN}_{NOTIFICATION_ID_2FA_LOGIN}"
-                not in notifications
-            ):
-                self._two_factor_url = None
+            self.hass, f"{DOMAIN}_{self.device.mac}_{notification_id}")
+
+    def _notification_dismiss_listener(self, type, data) -> None:
+        if type == persistent_notification.UpdateType.REMOVED and self.device:
+            notifications = self.hass.data.get(persistent_notification.DOMAIN)
+            if self._has_warning:
+                if (
+                    f"{DOMAIN}_{self.device.mac}_{NOTIFICATION_ID_WARNING}"
+                    not in notifications
+                ):
+                    if NOTIFICATION_ID_WARNING in self._notify:
+                        self.device.clear_warning()
+                    self._has_warning = self.device.status.has_warning
+            if self._two_factor_url:
+                if (
+                    f"{DOMAIN}_{self.device.mac}_{NOTIFICATION_ID_2FA_LOGIN}"
+                    not in notifications
+                ):
+                    self._two_factor_url = None
 
     def _fire_event(self, event_id, data) -> None:
         event_data =  {ATTR_ENTITY_ID: generate_entity_id("vacuum.{}", self.device.name, hass=self.hass)}
@@ -311,12 +322,18 @@ def _fire_event(self, event_id, data) -> None:
     async def _async_update_data(self) -> DreameVacuumDevice:
         """Handle device update. This function is only called once when the integration is added to Home Assistant."""
         try:
+            LOGGER.info("Integration starting...")
             await self.hass.async_add_executor_job(self.device.update)
             self.device.schedule_update()
             self.async_set_updated_data()
             return self.device
         except Exception as ex:
-            LOGGER.error("Update failed: %s", traceback.format_exc())
+            LOGGER.warning("Integration start failed: %s", traceback.format_exc())
+            if self.device is not None:
+                self.device.listen(None)
+                self.device.disconnect()
+                del self.device
+                self.device = None
             raise UpdateFailed(ex) from ex
 
     @callback

custom_components/dreame_vacuum/dreame/device.py

@@ -802,6 +802,7 @@ def disconnect(self) -> None:
         """Disconnect from device and cancel timers"""
         _LOGGER.info("Disconnect")
         self.schedule_update(-1)
+        self._protocol.disconnect()
         if self._map_manager:
             self._map_manager.schedule_update(-1)
 

custom_components/dreame_vacuum/dreame/map.py

@@ -171,6 +171,7 @@ def _init_data(self) -> None:
         self._map_request_time: int = None
         self._map_request_count: int = 0
         self._new_map_request_time: int = None
+        self._aes_iv: str = None
 
     def _request_map_from_cloud(self) -> bool:
         if self._current_timestamp_ms is not None:
@@ -599,7 +600,7 @@ def _get_interim_file_url(self, object_name: str) -> str | None:
         return url
 
     def _decode_map_partial(self, raw_map, timestamp=None, key=None) -> MapDataPartial | None:
-        partial_map = DreameVacuumMapDecoder.decode_map_partial(raw_map, key)
+        partial_map = DreameVacuumMapDecoder.decode_map_partial(raw_map, self._aes_iv, key)
         if partial_map is not None:
             # After restart or unsuccessful start robot returns timestamp_ms as uptime and that messes up with the latest map/frame id detection.
             # I could not figure out how app handles with this issue but i have added this code to update time stamp as request/object time.
@@ -1012,7 +1013,7 @@ def update(self) -> None:
 
     def set_aes_iv(self, aes_iv: str) -> None:
         if aes_iv:
-            DreameVacuumMapDecoder.AES_IV = aes_iv
+            self._aes_iv = aes_iv
 
     def set_vslam_map(self) -> None:
         self._vslam_map = True
@@ -1130,7 +1131,7 @@ def request_map_list(self) -> None:
                         if v.get(MAP_PARAMETER_MAP):
                             saved_map_data = DreameVacuumMapDecoder.decode_saved_map(
                                 v[MAP_PARAMETER_MAP], self._vslam_map, int(v[MAP_PARAMETER_ANGLE]) if v.get(
-                                    MAP_PARAMETER_ANGLE) else 0
+                                    MAP_PARAMETER_ANGLE) else 0, self._aes_iv
                             )
                             if saved_map_data is not None:
                                 name = v.get(MAP_PARAMETER_NAME)
@@ -1216,7 +1217,7 @@ def request_recovery_map_list(self) -> None:
                                     "Get recovery map file url result: %s", response)
                                 map_url = response[MAP_PARAMETER_RESULT][MAP_PARAMETER_URL]
                                 recovery_map_data = DreameVacuumMapDecoder.decode_saved_map(
-                                    map_info[MAP_PARAMETER_THB], self._vslam_map, self._saved_map_data[map_id].rotation)
+                                    map_info[MAP_PARAMETER_THB], self._vslam_map, self._saved_map_data[map_id].rotation, self._aes_iv)
                                 # TODO: store recovery map
 
     @property
@@ -1846,7 +1847,6 @@ def _current_timestamp_ms(self) -> int | None:
 
 class DreameVacuumMapDecoder:
     HEADER_SIZE = 27
-    AES_IV = ""
 
     @staticmethod
     def _read_int_8(data: bytes, offset: int = 0) -> int:
@@ -1985,7 +1985,7 @@ def _get_segment_center(map_data, segment_id: int, center: int, vertical: bool)
 
 
     @staticmethod
-    def decode_map_partial(raw_map, key=None) -> MapDataPartial | None:
+    def decode_map_partial(raw_map, iv=None, key=None) -> MapDataPartial | None:
         _LOGGER.debug("raw_map: %s", raw_map)
         raw_map = raw_map.replace("_", "/").replace("-", "+")
 
@@ -1997,12 +1997,13 @@ def decode_map_partial(raw_map, key=None) -> MapDataPartial | None:
         raw_map = base64.decodebytes(raw_map.encode("utf8"))
 
         if key is not None:
+            if iv is None:
+                iv = ""
             try:
                 key = hashlib.sha256(key.encode()).hexdigest()[
                     0:32].encode('utf8')
-                iv = DreameVacuumMapDecoder.AES_IV.encode('utf8')
                 cipher = Cipher(algorithms.AES(key), modes.CBC(
-                    iv), backend=default_backend())
+                    iv.encode("utf8")), backend=default_backend())
                 decryptor = cipher.decryptor()
                 raw_map = decryptor.update(raw_map) + decryptor.finalize()
             except Exception as ex:
@@ -2040,14 +2041,14 @@ def decode_map_partial(raw_map, key=None) -> MapDataPartial | None:
         return partial_map
 
     @staticmethod
-    def decode_map(raw_map: str, vslam_map: bool, rotation: int = 0) -> Tuple[MapData, Optional[MapData]]:
+    def decode_map(raw_map: str, vslam_map: bool, rotation: int = 0, iv: str = None, key: str = None) -> Tuple[MapData, Optional[MapData]]:
         return DreameVacuumMapDecoder.decode_map_data_from_partial(
-            DreameVacuumMapDecoder.decode_map_partial(raw_map), vslam_map, rotation
+            DreameVacuumMapDecoder.decode_map_partial(raw_map, iv, key), vslam_map, rotation
         )
 
     @staticmethod
-    def decode_saved_map(raw_map: str, vslam_map: bool, rotation: int = 0) -> MapData | None:
-        return DreameVacuumMapDecoder.decode_map(raw_map, vslam_map, rotation)[0]
+    def decode_saved_map(raw_map: str, vslam_map: bool, rotation: int = 0, iv: str = None) -> MapData | None:
+        return DreameVacuumMapDecoder.decode_map(raw_map, vslam_map, rotation, iv)[0]
 
     @staticmethod
     def decode_map_data_from_partial(

custom_components/dreame_vacuum/dreame/protocol.py

@@ -37,6 +37,9 @@ def set_credentials(self, ip: str, token: str):
     def connected(self) -> bool:
         return self._discovered
 
+    def disconnect(self):
+        self._discovered = False
+
 class DreameVacuumCloudProtocol:
     def __init__(self, username: str, password: str, country: str) -> None:
         self._username = username
@@ -355,6 +358,11 @@ def signed_nonce(self, nonce: str) -> str:
         )
         return base64.b64encode(hash_object.digest()).decode("utf-8")
 
+    def disconnect(self):
+        self._session.close()
+        self._connected = False
+        self._logged_in = False
+
     @staticmethod
     def generate_nonce():
         millis = int(round(time.time() * 1000))
@@ -489,6 +497,15 @@ def connect(self, retry_count=1) -> Any:
             self._connected = True
         return response
 
+    def disconnect(self):
+        if self.device is not None:
+            self.device.disconnect()
+        if self.cloud is not None:
+            self.cloud.disconnect()
+        if self.device_cloud is not None:
+            self.device_cloud.disconnect()
+        self._connected = False
+
     def send(self, method, parameters: Any = None, retry_count: int = 1) -> Any:
         if (self.prefer_cloud or not self.device) and self.device_cloud:
             if not self.device_cloud.logged_in:

custom_components/dreame_vacuum/manifest.json

@@ -16,7 +16,8 @@
     "pycryptodome",
     "python-miio",
     "py-mini-racer",
-    "tzlocal"
+    "tzlocal",
+    "paho-mqtt"
   ],
-  "version": "v1.0.1"
+  "version": "v1.0.2"
 }