From 1e2312087238eb22f9aa03e59ac5336b4dc0054c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pierre=20St=C3=A5hl?= Date: Thu, 17 Sep 2020 13:37:33 +0200 Subject: [PATCH] Add device discovery to config flow --- custom_components/localtuya/config_flow.py | 75 ++++++++++++++--- custom_components/localtuya/discovery.py | 81 +++++++++++++++++++ .../localtuya/translations/en.json | 12 ++- 3 files changed, 156 insertions(+), 12 deletions(-) create mode 100644 custom_components/localtuya/discovery.py diff --git a/custom_components/localtuya/config_flow.py b/custom_components/localtuya/config_flow.py index 12dcac7..6368ed4 100644 --- a/custom_components/localtuya/config_flow.py +++ b/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={}) diff --git a/custom_components/localtuya/discovery.py b/custom_components/localtuya/discovery.py new file mode 100644 index 0000000..bca785f --- /dev/null +++ b/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() diff --git a/custom_components/localtuya/translations/en.json b/custom_components/localtuya/translations/en.json index 8bfbe58..9ef9c65 100644 --- a/custom_components/localtuya/translations/en.json +++ b/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",