Add device discovery to config flow

This commit is contained in:
Pierre Ståhl
2020-09-17 13:37:33 +02:00
committed by rospogrigio
parent ffae9976f0
commit 1e23120872
3 changed files with 156 additions and 12 deletions

View File

@@ -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()
@@ -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={})

View File

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

View File

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