Add device discovery to config flow

custom_components/localtuya/config_flow.py

@@ -1,14 +1,12 @@
 """Config flow for LocalTuya integration integration."""
 import logging
 from importlib import import_module
-from itertools import chain
 
 import voluptuous as vol
 
 from homeassistant import config_entries, core, exceptions
 from homeassistant.core import callback
 from homeassistant.const import (
-    CONF_NAME,
     CONF_ENTITIES,
     CONF_ID,
     CONF_HOST,
@@ -27,13 +25,19 @@ from .const import (  # pylint: disable=unused-import
     DOMAIN,
     PLATFORMS,
 )
+from .discovery import discover
 
 _LOGGER = logging.getLogger(__name__)
 
+DISCOVER_TIMEOUT = 6.0
+
 PLATFORM_TO_ADD = "platform_to_add"
 NO_ADDITIONAL_PLATFORMS = "no_additional_platforms"
+DISCOVERED_DEVICE = "discovered_device"
 
-USER_SCHEMA = vol.Schema(
+CUSTOM_DEVICE = "..."
+
+BASIC_INFO_SCHEMA = vol.Schema(
     {
         vol.Required(CONF_FRIENDLY_NAME): str,
         vol.Required(CONF_HOST): str,
@@ -57,6 +61,14 @@ PICK_ENTITY_SCHEMA = vol.Schema(
 )
 
 
+def user_schema(devices):
+    """Create schema for user step."""
+    devices = [f"{ip} ({dev['gwId']})" for ip, dev in devices.items()]
+    return vol.Schema(
+        {vol.Required(DISCOVERED_DEVICE): vol.In(devices + [CUSTOM_DEVICE])}
+    )
+
+
 def schema_defaults(schema, dps_list=None, **defaults):
     """Create a new schema with default values filled in."""
     copy = schema.extend({})
@@ -125,7 +137,9 @@ async def validate_input(hass: core.HomeAssistant, data):
     detected_dps = {}
 
     try:
-        detected_dps = await hass.async_add_executor_job(pytuyadevice.detect_available_dps)
+        detected_dps = await hass.async_add_executor_job(
+            pytuyadevice.detect_available_dps
+        )
     except (ConnectionRefusedError, ConnectionResetError):
         raise CannotConnect
     except ValueError:
@@ -151,10 +165,35 @@ class LocaltuyaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         self.basic_info = None
         self.dps_strings = []
         self.platform = None
+        self.devices = {}
+        self.selected_device = None
         self.entities = []
 
     async def async_step_user(self, user_input=None):
-        """Handle the initial step."""
+        """Handle initial step."""
+        errors = {}
+        if user_input is not None:
+            if user_input[DISCOVERED_DEVICE] != CUSTOM_DEVICE:
+                self.selected_device = user_input[DISCOVERED_DEVICE].split(" ")[0]
+            return await self.async_step_basic_info()
+
+        try:
+            devices = await discover(DISCOVER_TIMEOUT, self.hass.loop)
+            self.devices = {
+                ip: dev
+                for ip, dev in devices.items()
+                if dev["gwId"] not in self._async_current_ids()
+            }
+        except Exception:  # pylint: disable=broad-except
+            _LOGGER.exception("discovery failed")
+            errors["base"] = "discovery_failed"
+
+        return self.async_show_form(
+            step_id="user", errors=errors, data_schema=user_schema(self.devices)
+        )
+
+    async def async_step_basic_info(self, user_input=None):
+        """Handle input of basic info."""
         errors = {}
         if user_input is not None:
             await self.async_set_unique_id(user_input[CONF_DEVICE_ID])
@@ -172,8 +211,18 @@ class LocaltuyaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                 _LOGGER.exception("Unexpected exception")
                 errors["base"] = "unknown"
 
+        defaults = {}
+        defaults.update(user_input or {})
+        if self.selected_device is not None:
+            device = self.devices[self.selected_device]
+            defaults[CONF_HOST] = device.get("ip")
+            defaults[CONF_DEVICE_ID] = device.get("gwId")
+            defaults[CONF_PROTOCOL_VERSION] = device.get("version")
+
         return self.async_show_form(
-            step_id="user", data_schema=USER_SCHEMA, errors=errors
+            step_id="basic_info",
+            data_schema=schema_defaults(BASIC_INFO_SCHEMA, **defaults),
+            errors=errors,
         )
 
     async def async_step_pick_entity_type(self, user_input=None):
@@ -185,7 +234,9 @@ class LocaltuyaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                     CONF_DPS_STRINGS: self.dps_strings,
                     CONF_ENTITIES: self.entities,
                 }
-                return self.async_create_entry(title=config[CONF_FRIENDLY_NAME], data=config)
+                return self.async_create_entry(
+                    title=config[CONF_FRIENDLY_NAME], data=config
+                )
 
             self.platform = user_input[PLATFORM_TO_ADD]
             return await self.async_step_add_entity()
@@ -243,7 +294,7 @@ class LocaltuyaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         else:
             self.entities.append(_convert_entity(user_input))
 
-        #print('ENTITIES: [{}] '.format(self.entities))
+        # print('ENTITIES: [{}] '.format(self.entities))
         config = {
             CONF_FRIENDLY_NAME: f"{user_input[CONF_FRIENDLY_NAME]}",
             CONF_HOST: user_input[CONF_HOST],
@@ -254,7 +305,9 @@ class LocaltuyaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
             CONF_ENTITIES: self.entities,
         }
         self._abort_if_unique_id_configured(updates=config)
-        return self.async_create_entry(title=f"{config[CONF_FRIENDLY_NAME]} (YAML)", data=config)
+        return self.async_create_entry(
+            title=f"{config[CONF_FRIENDLY_NAME]} (YAML)", data=config
+        )
 
 
 class LocalTuyaOptionsFlowHandler(config_entries.OptionsFlow):
@@ -300,7 +353,9 @@ class LocalTuyaOptionsFlowHandler(config_entries.OptionsFlow):
 
             if len(self.entities) == len(self.data[CONF_ENTITIES]):
                 self.hass.config_entries.async_update_entry(
-                    self.config_entry, title=self.data[CONF_FRIENDLY_NAME], data=self.data
+                    self.config_entry,
+                    title=self.data[CONF_FRIENDLY_NAME],
+                    data=self.data,
                 )
                 return self.async_create_entry(title="", data={})
 

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()

custom_components/localtuya/translations/en.json

@@ -7,12 +7,20 @@
             "cannot_connect": "Cannot connect to device. Verify that address is correct and try again.",
             "invalid_auth": "Failed to authenticate with device. Verify that device id and local key are correct.",
             "unknown": "An unknown error occurred. See log for details.",
-            "entity_already_configured": "Entity with this ID has already been configured."
+            "entity_already_configured": "Entity with this ID has already been configured.",
+            "discovery_failed": "Failed to discover devices. You can still add a device manually."
         },
         "step": {
             "user": {
+                "title": "Device Discovery",
+                "description": "Pick one of the automatically discovered devices or `...` to manually to add a device.",
+                "data": {
+                    "discovered_device": "Discovered Device"
+                }
+            },
+            "basic_info": {
                 "title": "Add Tuya device",
-                "description": "Fill in the basic details and pick device type. The name entered here will be used to identify the integration itself (as seen in the `Integrations` page). You will add entities and give them names in the following steps.",
+                "description": "Fill in the basic device details. The name entered here will be used to identify the integration itself (as seen in the `Integrations` page). You will add entities and give them names in the following steps.",
                 "data": {
                     "friendly_name": "Name",
                     "host": "Host",