Add device discovery to config flow

custom_components/localtuya/discovery.py

@@ -0,0 +1,81 @@
+"""Discovery module for Tuya devices.
+
+Entirely based on tuya-convert.py from tuya-convert:
+
+https://github.com/ct-Open-Source/tuya-convert/blob/master/scripts/tuya-discovery.py
+"""
+import json
+import asyncio
+import logging
+from hashlib import md5
+
+from Cryptodome.Cipher import AES
+
+_LOGGER = logging.getLogger(__name__)
+
+UDP_KEY = md5(b"yGAdlopoPVldABfn").digest()
+
+
+def decrypt_udp(message):
+    """Decrypt encrypted UDP broadcasts."""
+    def _unpad(data):
+        return data[:-ord(data[len(data) - 1:])]
+
+    return _unpad(AES.new(UDP_KEY, AES.MODE_ECB).decrypt(message)).decode()
+
+
+class TuyaDiscovery(asyncio.DatagramProtocol):
+    """Datagram handler listening for Tuya broadcast messages."""
+
+    def __init__(self, found_devices):
+        """Initialize a new TuyaDiscovery instance."""
+        self.found_devices = found_devices
+
+    def datagram_received(self, data, addr):
+        """Handle received broadcast message."""
+        data = data[20:-8]
+        try:
+            data = decrypt_udp(data)
+        except Exception:
+            data = data.decode()
+
+        decoded = json.loads(data)
+        if decoded.get("ip") not in self.found_devices:
+            self.found_devices[decoded.get("ip")] = decoded
+            _LOGGER.debug("Discovered device: %s", decoded)
+
+
+async def discover(timeout, loop):
+    """Discover and return Tuya devices on the network."""
+    found_devices = {}
+
+    def proto_factory():
+        return TuyaDiscovery(found_devices)
+
+    listener = loop.create_datagram_endpoint(
+        proto_factory, local_addr=("0.0.0.0", 6666)
+    )
+    encrypted_listener = loop.create_datagram_endpoint(
+        proto_factory, local_addr=("0.0.0.0", 6667)
+    )
+
+    listeners = await asyncio.gather(listener, encrypted_listener)
+    _LOGGER.debug("Listening to broadcasts on UDP port 6666 and 6667")
+
+    try:
+        await asyncio.sleep(timeout)
+    finally:
+        for transport, _ in listeners:
+            transport.close()
+
+    return found_devices
+
+
+def main():
+    loop = asyncio.get_event_loop()
+    res = loop.run_until_complete(discover(5, loop))
+    print(res)
+
+
+if __name__ == "__main__":
+    main()