diff --git a/custom_components/localtuya/__init__.py b/custom_components/localtuya/__init__.py index be6a8db..f60908d 100644 --- a/custom_components/localtuya/__init__.py +++ b/custom_components/localtuya/__init__.py @@ -63,13 +63,16 @@ from homeassistant.const import ( CONF_DEVICE_ID, CONF_PLATFORM, CONF_ENTITIES, + CONF_HOST, + EVENT_HOMEASSISTANT_STOP, SERVICE_RELOAD, ) from homeassistant.helpers.reload import async_integration_yaml_config -from .const import DOMAIN, TUYA_DEVICE +from .const import DATA_DISCOVERY, DOMAIN, TUYA_DEVICE from .config_flow import config_schema from .common import TuyaDevice +from .discovery import TuyaDiscovery _LOGGER = logging.getLogger(__name__) @@ -92,6 +95,8 @@ async def async_setup(hass: HomeAssistant, config: dict): """Set up the LocalTuya integration component.""" hass.data.setdefault(DOMAIN, {}) + device_cache = {} + async def _handle_reload(service): """Handle reload service call.""" config = await async_integration_yaml_config(hass, DOMAIN) @@ -112,6 +117,51 @@ async def async_setup(hass: HomeAssistant, config: dict): await asyncio.gather(*reload_tasks) + def _entry_by_device_id(device_id): + """Look up config entry by device id.""" + current_entries = hass.config_entries.async_entries(DOMAIN) + for entry in current_entries: + if entry.data[CONF_DEVICE_ID] == device_id: + return entry + return None + + def _device_discovered(device): + """Update address of device if it has changed.""" + device_ip = device["ip"] + device_id = device["gwId"] + + # If device is not in cache, check if a config entry exists + if device_id not in device_cache: + entry = _entry_by_device_id(device_id) + if entry: + # Save address from config entry in cache to trigger + # potential update below + device_cache[device_id] = entry.data[CONF_HOST] + + # If device is in cache and address changed... + if device_id in device_cache and device_cache[device_id] != device_ip: + _LOGGER.debug("Device %s changed IP to %s", device_id, device_ip) + + entry = _entry_by_device_id(device_id) + if entry: + hass.config_entries.async_update_entry( + entry, data={**entry.data, CONF_HOST: device_ip} + ) + device_cache[device_id] = device_ip + + discovery = TuyaDiscovery(_device_discovered) + + def _shutdown(event): + """Clean up resources when shutting down.""" + discovery.close() + + try: + await discovery.start() + hass.data[DOMAIN][DATA_DISCOVERY] = discovery + hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _shutdown) + except Exception: + _LOGGER.exception("failed to set up discovery") + hass.helpers.service.async_register_admin_service( DOMAIN, SERVICE_RELOAD, diff --git a/custom_components/localtuya/config_flow.py b/custom_components/localtuya/config_flow.py index c86bd6a..5247c10 100644 --- a/custom_components/localtuya/config_flow.py +++ b/custom_components/localtuya/config_flow.py @@ -21,15 +21,13 @@ from .const import ( # pylint: disable=unused-import CONF_LOCAL_KEY, CONF_PROTOCOL_VERSION, CONF_DPS_STRINGS, + DATA_DISCOVERY, 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" @@ -211,16 +209,12 @@ class LocaltuyaConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): 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" + devices = self.hass.data[DOMAIN][DATA_DISCOVERY].devices + self.devices = { + ip: dev + for ip, dev in devices.items() + if dev["gwId"] not in self._async_current_ids() + } return self.async_show_form( step_id="user", errors=errors, data_schema=user_schema(self.devices) diff --git a/custom_components/localtuya/const.py b/custom_components/localtuya/const.py index 4a6195b..afd6f23 100644 --- a/custom_components/localtuya/const.py +++ b/custom_components/localtuya/const.py @@ -28,6 +28,8 @@ CONF_SPAN_TIME = "span_time" # sensor CONF_SCALING = "scaling" +DATA_DISCOVERY = "discovery" + DOMAIN = "localtuya" # Platforms in this list must support config flows diff --git a/custom_components/localtuya/discovery.py b/custom_components/localtuya/discovery.py index 5a64659..a75cc22 100644 --- a/custom_components/localtuya/discovery.py +++ b/custom_components/localtuya/discovery.py @@ -31,9 +31,29 @@ def decrypt_udp(message): 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 __init__(self, callback): + """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() + listener = loop.create_datagram_endpoint( + lambda: self, local_addr=("0.0.0.0", 6666) + ) + encrypted_listener = loop.create_datagram_endpoint( + lambda: self, local_addr=("0.0.0.0", 6667) + ) + + self.listeners = await asyncio.gather(listener, encrypted_listener) + _LOGGER.debug("Listening to broadcasts on UDP port 6666 and 6667") + + def close(self): + """Stop discovery.""" + for transport, _ in self.listeners: + transport.close() def datagram_received(self, data, addr): """Handle received broadcast message.""" @@ -44,43 +64,12 @@ class TuyaDiscovery(asyncio.DatagramProtocol): 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) + self.device_found(decoded) + def device_found(self, device): + """Discover a new device.""" + if device.get("ip") not in self.devices: + self.devices[device.get("ip")] = device + _LOGGER.debug("Discovered device: %s", device) -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(): - """Run discovery and print result.""" - loop = asyncio.get_event_loop() - res = loop.run_until_complete(discover(5, loop)) - print(res) - - -if __name__ == "__main__": - main() + self._callback(device) diff --git a/custom_components/localtuya/translations/en.json b/custom_components/localtuya/translations/en.json index d0b4d45..536350f 100644 --- a/custom_components/localtuya/translations/en.json +++ b/custom_components/localtuya/translations/en.json @@ -7,8 +7,7 @@ "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.", - "discovery_failed": "Failed to discover devices. You can still add a device manually." + "entity_already_configured": "Entity with this ID has already been configured." }, "step": { "user": {