"""Config flow for TP-Link.""" from __future__ import annotations from typing import Any from kasa import SmartDevice, SmartDeviceException from kasa.discover import Discover import voluptuous as vol from homeassistant import config_entries from homeassistant.components import dhcp from homeassistant.const import CONF_DEVICE, CONF_HOST, CONF_MAC from homeassistant.core import callback from homeassistant.data_entry_flow import FlowResult from homeassistant.helpers import device_registry as dr from homeassistant.helpers.typing import DiscoveryInfoType from . import async_discover_devices from .const import DOMAIN class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): """Handle a config flow for tplink.""" VERSION = 1 def __init__(self) -> None: """Initialize the config flow.""" self._discovered_devices: dict[str, SmartDevice] = {} self._discovered_device: SmartDevice | None = None async def async_step_dhcp(self, discovery_info: dhcp.DhcpServiceInfo) -> FlowResult: """Handle discovery via dhcp.""" return await self._async_handle_discovery( discovery_info.ip, discovery_info.macaddress ) async def async_step_integration_discovery( self, discovery_info: DiscoveryInfoType ) -> FlowResult: """Handle integration discovery.""" return await self._async_handle_discovery( discovery_info[CONF_HOST], discovery_info[CONF_MAC] ) async def _async_handle_discovery(self, host: str, mac: str) -> FlowResult: """Handle any discovery.""" await self.async_set_unique_id(dr.format_mac(mac)) self._abort_if_unique_id_configured(updates={CONF_HOST: host}) self._async_abort_entries_match({CONF_HOST: host}) self.context[CONF_HOST] = host for progress in self._async_in_progress(): if progress.get("context", {}).get(CONF_HOST) == host: return self.async_abort(reason="already_in_progress") try: self._discovered_device = await self._async_try_connect( host, raise_on_progress=True ) except SmartDeviceException: return self.async_abort(reason="cannot_connect") return await self.async_step_discovery_confirm() async def async_step_discovery_confirm( self, user_input: dict[str, Any] | None = None ) -> FlowResult: """Confirm discovery.""" assert self._discovered_device is not None if user_input is not None: return self._async_create_entry_from_device(self._discovered_device) self._set_confirm_only() placeholders = { "name": self._discovered_device.alias, "model": self._discovered_device.model, "host": self._discovered_device.host, } self.context["title_placeholders"] = placeholders return self.async_show_form( step_id="discovery_confirm", description_placeholders=placeholders ) async def async_step_user( self, user_input: dict[str, Any] | None = None ) -> FlowResult: """Handle the initial step.""" errors = {} if user_input is not None: if not (host := user_input[CONF_HOST]): return await self.async_step_pick_device() try: device = await self._async_try_connect(host, raise_on_progress=False) except SmartDeviceException: errors["base"] = "cannot_connect" else: return self._async_create_entry_from_device(device) return self.async_show_form( step_id="user", data_schema=vol.Schema({vol.Optional(CONF_HOST, default=""): str}), errors=errors, ) async def async_step_pick_device( self, user_input: dict[str, Any] | None = None ) -> FlowResult: """Handle the step to pick discovered device.""" if user_input is not None: mac = user_input[CONF_DEVICE] await self.async_set_unique_id(mac, raise_on_progress=False) return self._async_create_entry_from_device(self._discovered_devices[mac]) configured_devices = { entry.unique_id for entry in self._async_current_entries() } self._discovered_devices = await async_discover_devices(self.hass) devices_name = { formatted_mac: f"{device.alias} {device.model} ({device.host}) {formatted_mac}" for formatted_mac, device in self._discovered_devices.items() if formatted_mac not in configured_devices } # Check if there is at least one device if not devices_name: return self.async_abort(reason="no_devices_found") return self.async_show_form( step_id="pick_device", data_schema=vol.Schema({vol.Required(CONF_DEVICE): vol.In(devices_name)}), ) @callback def _async_create_entry_from_device(self, device: SmartDevice) -> FlowResult: """Create a config entry from a smart device.""" self._abort_if_unique_id_configured(updates={CONF_HOST: device.host}) return self.async_create_entry( title=f"{device.alias} {device.model}", data={ CONF_HOST: device.host, }, ) async def _async_try_connect( self, host: str, raise_on_progress: bool = True ) -> SmartDevice: """Try to connect.""" self._async_abort_entries_match({CONF_HOST: host}) device: SmartDevice = await Discover.discover_single(host) await self.async_set_unique_id( dr.format_mac(device.mac), raise_on_progress=raise_on_progress ) return device