Files
rikku-home-assistant/custom_components/localtuya/discovery.py
T
2025-04-21 21:03:30 -04:00

131 lines
4.2 KiB
Python

"""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