Add local discovery of devices

custom_components/robovac/__init__.py

@@ -15,39 +15,60 @@
 
 """The Eufy Robovac integration."""
 from __future__ import annotations
+import logging
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import Platform
+from homeassistant.const import EVENT_HOMEASSISTANT_STOP, Platform
 from homeassistant.core import HomeAssistant
-from .const import DOMAIN
+from .const import CONF_VACS, DOMAIN
 
-PLATFORMS = [Platform.VACUUM]
+from .tuyalocaldiscovery import TuyaLocalDiscovery
+
+PLATFORM = Platform.VACUUM
+_LOGGER = logging.getLogger(__name__)
+
+
+async def async_setup(hass, entry) -> bool:
+    hass.data.setdefault(DOMAIN, {})
+
+    current_entries = hass.config_entries.async_entries(DOMAIN)
+    hass_data = dict(current_entries[0].data)
+
+    def update_device(device):
+        if device["gwId"] in hass_data[CONF_VACS]:
+            if hass_data[CONF_VACS][device["gwId"]]["ip_address"] != device.ip:
+                hass_data[CONF_VACS][device["gwId"]]["ip_address"] = device.ip
+                hass.config_entries.async_update_entry(entry, data=hass_data)
+
+    tuyalocaldiscovery = TuyaLocalDiscovery(update_device)
+    try:
+        await tuyalocaldiscovery.start()
+        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, tuyalocaldiscovery.close)
+    except Exception:  # pylint: disable=broad-except
+        _LOGGER.exception("failed to set up discovery")
+
+    return True
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Eufy Robovac from a config entry."""
-    hass.data.setdefault(DOMAIN, {})
 
-    #    hass_data = dict(entry.data)
-    # Registers update listener to update config entry when options are updated.
-    #    unsub_options_update_listener = entry.add_update_listener(options_update_listener)
-    # Store a reference to the unsubscribe function to cleanup if an entry is unloaded.
-    #    hass_data["unsub_options_update_listener"] = unsub_options_update_listener
-    #    hass.data[DOMAIN][entry.entry_id] = hass_data
     entry.async_on_unload(entry.add_update_listener(update_listener))
 
-    await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
+    await hass.config_entries.async_forward_entry_setup(entry, PLATFORM)
+
     return True
 
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
+    if unload_ok := await hass.config_entries.async_forward_entry_unload(
+        entry, PLATFORM
+    ):
         """Nothing"""
     return unload_ok
 
 
 async def update_listener(hass, entry):
     """Handle options update."""
-    await async_unload_entry(hass, entry)
-    await async_setup_entry(hass, entry)
+    await hass.config_entries.async_reload(entry.entry_id)

custom_components/robovac/tuyalocaldiscovery.py

@@ -0,0 +1,57 @@
+import asyncio
+import json
+import logging
+from hashlib import md5
+
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+
+_LOGGER = logging.getLogger(__name__)
+
+UDP_KEY = md5(b"yGAdlopoPVldABfn").digest()
+
+class TuyaLocalDiscovery(asyncio.DatagramProtocol):
+    def __init__(self, callback):
+        self.devices = {}
+        self._listeners = []
+        self.discovered_callback = callback
+
+    async def start(self):
+        loop = asyncio.get_running_loop()
+        listener = loop.create_datagram_endpoint(
+            lambda: self, local_addr=("0.0.0.0", 6666)
+        )
+        encrypted_listener = loop.create_datagram_endpoint(
+            lambda: self, local_addr=("0.0.0.0", 6667)
+        )
+
+        self._listeners = await asyncio.gather(listener, encrypted_listener)
+        _LOGGER.debug("Listening to broadcasts on UDP port 6666 and 6667")
+
+    def close(self):
+        for transport, _ in self._listeners:
+            transport.close()
+
+    def datagram_received(self, data, addr):
+        data = data[20:-8]
+        try:
+            cipher = Cipher(algorithms.AES(UDP_KEY), modes.ECB(), default_backend())
+            decryptor = cipher.decryptor()
+            padded_data = decryptor.update(data) + decryptor.finalize()
+            data = padded_data[: -ord(data[len(data) - 1 :])]
+
+        except Exception:
+            data = data.decode()
+
+        decoded = json.loads(data)
+        self.discovered_callback(decoded)
+
+async def discover():
+    """Discover and return devices on local network."""
+    discovery = TuyaDiscovery()
+    try:
+        await discovery.start()
+        await asyncio.sleep(DEFAULT_TIMEOUT)
+    finally:
+        discovery.close()
+    return discovery.devices