Initial commit.

custom_components/localtuya/discovery.py

@@ -0,0 +1,130 @@
+"""Discovery module for Tuya devices.
+
+based on tuya-convert.py from tuya-convert:
+    https://github.com/ct-Open-Source/tuya-convert/blob/master/scripts/tuya-discovery.py
+
+Maintained by @xZetsubou
+"""
+
+import os
+import asyncio
+import json
+import logging
+from hashlib import md5
+from socket import inet_aton
+
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+
+from .entity import pytuya
+
+_LOGGER = logging.getLogger(__name__)
+
+UDP_KEY = md5(b"yGAdlopoPVldABfn").digest()
+
+PREFIX_55AA_BIN = b"\x00\x00U\xaa"
+PREFIX_6699_BIN = b"\x00\x00\x66\x99"
+UDP_COMMAND = b"\x00\x00\x00\x00"
+
+DEFAULT_TIMEOUT = 6.0
+
+
+def decrypt(msg, key):
+    def _unpad(data):
+        return data[: -ord(data[len(data) - 1 :])]
+
+    cipher = Cipher(algorithms.AES(key), modes.ECB(), default_backend())
+    decryptor = cipher.decryptor()
+    return _unpad(decryptor.update(msg) + decryptor.finalize()).decode()
+
+
+def decrypt_udp(message):
+    """Decrypt encrypted UDP broadcasts."""
+    if message[:4] == PREFIX_55AA_BIN:
+        payload = message[20:-8]
+        if message[8:12] == UDP_COMMAND:
+            return payload
+        return decrypt(payload, UDP_KEY)
+    if message[:4] == PREFIX_6699_BIN:
+        unpacked = pytuya.unpack_message(message, hmac_key=UDP_KEY, no_retcode=None)
+        payload = unpacked.payload.decode()
+        # app sometimes has extra bytes at the end
+        while payload[-1] == chr(0):
+            payload = payload[:-1]
+        return payload
+    return decrypt(message, UDP_KEY)
+
+
+class TuyaDiscovery(asyncio.DatagramProtocol):
+    """Datagram handler listening for Tuya broadcast messages."""
+
+    def __init__(self, callback=None):
+        """Initialize a new BaseDiscovery."""
+        self.devices = {}
+        self._listeners = []
+        self._callback = callback
+
+    async def start(self):
+        """Start discovery by listening to broadcasts."""
+        loop = asyncio.get_running_loop()
+        op_reuse_port = {"reuse_port": True} if os.name != "nt" else {}
+        listener = loop.create_datagram_endpoint(
+            lambda: self, local_addr=("0.0.0.0", 6666), **op_reuse_port
+        )
+        encrypted_listener = loop.create_datagram_endpoint(
+            lambda: self, local_addr=("0.0.0.0", 6667), **op_reuse_port
+        )
+        # tuyaApp_encrypted_listener = loop.create_datagram_endpoint(
+        #     lambda: self, local_addr=("0.0.0.0", 7000), **op_reuse_port
+        # )
+        self._listeners = await asyncio.gather(listener, encrypted_listener)
+        _LOGGER.debug("Listening to broadcasts on UDP port 6666, 6667")
+
+    def close(self):
+        """Stop discovery."""
+        self._callback = None
+        for transport, _ in self._listeners:
+            transport.close()
+
+    def datagram_received(self, data, addr):
+        """Handle received broadcast message."""
+        try:
+            try:
+                data = decrypt_udp(data)
+            except Exception:  # pylint: disable=broad-except
+                data = data.decode()
+            decoded = json.loads(data)
+            self.device_found(decoded)
+        except:
+            # _LOGGER.debug("Bordcast from app from ip: %s", addr[0])
+            _LOGGER.debug("Failed to decode broadcast from %r: %r", addr[0], data)
+
+    def device_found(self, device):
+        """Discover a new device."""
+        gwid, ip = device.get("gwId"), device.get("ip")
+        # If device found but the ip changed.
+        if gwid in self.devices and (self.devices[gwid].get("ip") != ip):
+            self.devices.pop(gwid)
+
+        if gwid not in self.devices:
+            self.devices[gwid] = device
+            # Sort devices by ip.
+            sort_devices = sorted(
+                self.devices.items(), key=lambda i: inet_aton(i[1].get("ip", "0"))
+            )
+            self.devices = dict(sort_devices)
+
+            _LOGGER.debug("Discovered device: %s", device)
+        if self._callback:
+            self._callback(device)
+
+
+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