Improve UPnP configuration flow

homeassistant/components/discovery/__init__.py

@@ -32,7 +32,6 @@
 SERVICE_HASS_IOS_APP = "hass_ios"
 SERVICE_HASSIO = "hassio"
 SERVICE_HEOS = "heos"
-SERVICE_IGD = "igd"
 SERVICE_KONNECTED = "konnected"
 SERVICE_MOBILE_APP = "hass_mobile_app"
 SERVICE_NETGEAR = "netgear_router"
@@ -48,7 +47,6 @@
 CONFIG_ENTRY_HANDLERS = {
     SERVICE_DAIKIN: "daikin",
     SERVICE_TELLDUSLIVE: "tellduslive",
-    SERVICE_IGD: "upnp",
 }
 
 SERVICE_HANDLERS = {

homeassistant/components/upnp/__init__.py

@@ -19,6 +19,12 @@
     CONF_HASS,
     CONF_LOCAL_IP,
     CONF_PORTS,
+    CONFIG_ENTRY_ST,
+    CONFIG_ENTRY_UDN,
+    DISCOVERY_LOCATION,
+    DISCOVERY_ST,
+    DISCOVERY_UDN,
+    DISCOVERY_USN,
     DOMAIN,
     LOGGER as _LOGGER,
 )
@@ -89,40 +95,41 @@ async def async_discover_and_construct(
     """Discovery devices and construct a Device for one."""
     # pylint: disable=invalid-name
     discovery_infos = await Device.async_discover(hass)
+    _LOGGER.debug("Discovered devices: %s", discovery_infos)
     if not discovery_infos:
         _LOGGER.info("No UPnP/IGD devices discovered")
         return None
 
     if udn:
-        # get the discovery info with specified UDN
-        _LOGGER.debug("Discovery_infos: %s", discovery_infos)
-        filtered = [di for di in discovery_infos if di["udn"] == udn]
+        # Get the discovery info with specified UDN/ST.
+        filtered = [di for di in discovery_infos if di[DISCOVERY_UDN] == udn]
         if st:
-            _LOGGER.debug("Filtering on ST: %s", st)
-            filtered = [di for di in discovery_infos if di["st"] == st]
+            filtered = [di for di in discovery_infos if di[DISCOVERY_ST] == st]
         if not filtered:
             _LOGGER.warning(
-                'Wanted UPnP/IGD device with UDN "%s" not found, ' "aborting", udn
+                'Wanted UPnP/IGD device with UDN "%s" not found, aborting', udn
             )
             return None
-        # ensure we're always taking the latest
-        filtered = sorted(filtered, key=itemgetter("st"), reverse=True)
+
+        # Ensure we're always taking the latest, if we filtered only on UDN.
+        filtered = sorted(filtered, key=itemgetter(DISCOVERY_ST), reverse=True)
         discovery_info = filtered[0]
     else:
-        # get the first/any
+        # Get the first/any.
         discovery_info = discovery_infos[0]
         if len(discovery_infos) > 1:
             device_name = discovery_info.get(
-                "usn", discovery_info.get("ssdp_description", "")
+                DISCOVERY_USN, discovery_info.get(DISCOVERY_LOCATION, "")
             )
             _LOGGER.info("Detected multiple UPnP/IGD devices, using: %s", device_name)
 
-    ssdp_description = discovery_info["ssdp_description"]
-    return await Device.async_create_device(hass, ssdp_description)
+    location = discovery_info[DISCOVERY_LOCATION]
+    return await Device.async_create_device(hass, location)
 
 
 async def async_setup(hass: HomeAssistantType, config: ConfigType):
     """Set up UPnP component."""
+    _LOGGER.debug("async_setup, config: %s", config)
     conf_default = CONFIG_SCHEMA({DOMAIN: {}})[DOMAIN]
     conf = config.get(DOMAIN, conf_default)
     local_ip = await hass.async_add_executor_job(get_local_ip)
@@ -133,7 +140,8 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType):
         "ports": conf.get(CONF_PORTS),
     }
 
-    if conf is not None:
+    # Only start if set up via configuration.yaml.
+    if DOMAIN in config:
         hass.async_create_task(
             hass.config_entries.flow.async_init(
                 DOMAIN, context={"source": config_entries.SOURCE_IMPORT}
@@ -145,23 +153,26 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType):
 
 async def async_setup_entry(hass: HomeAssistantType, config_entry: ConfigEntry) -> bool:
     """Set up UPnP/IGD device from a config entry."""
+    _LOGGER.debug("async_setup_entry, config_entry: %s", config_entry.data)
     domain_data = hass.data[DOMAIN]
     conf = domain_data["config"]
 
     # discover and construct
-    udn = config_entry.data.get("udn")
-    st = config_entry.data.get("st")  # pylint: disable=invalid-name
+    udn = config_entry.data.get(CONFIG_ENTRY_UDN)
+    st = config_entry.data.get(CONFIG_ENTRY_ST)  # pylint: disable=invalid-name
     device = await async_discover_and_construct(hass, udn, st)
     if not device:
         _LOGGER.info("Unable to create UPnP/IGD, aborting")
         raise ConfigEntryNotReady
 
-    # 'register'/save UDN + ST
+    # 'register'/save device
     hass.data[DOMAIN]["devices"][device.udn] = device
-    hass.config_entries.async_update_entry(
-        entry=config_entry,
-        data={**config_entry.data, "udn": device.udn, "st": device.device_type},
-    )
+
+    # Ensure entry has proper unique_id.
+    if config_entry.unique_id != device.unique_id:
+        hass.config_entries.async_update_entry(
+            entry=config_entry, unique_id=device.unique_id,
+        )
 
     # create device registry entry
     device_registry = await dr.async_get_registry(hass)
@@ -211,7 +222,7 @@ async def async_unload_entry(
     hass: HomeAssistantType, config_entry: ConfigEntry
 ) -> bool:
     """Unload a UPnP/IGD device from a config entry."""
-    udn = config_entry.data["udn"]
+    udn = config_entry.data[CONFIG_ENTRY_UDN]
     device = hass.data[DOMAIN]["devices"][udn]
 
     # remove port mapping

homeassistant/components/upnp/config_flow.py

@@ -1,10 +1,187 @@
 """Config flow for UPNP."""
+from typing import Mapping, Optional
+
+import voluptuous as vol
+
 from homeassistant import config_entries
-from homeassistant.helpers import config_entry_flow
+from homeassistant.components import ssdp
 
-from .const import DOMAIN
+from .const import (  # pylint: disable=unused-import
+    CONFIG_ENTRY_ST,
+    CONFIG_ENTRY_UDN,
+    DISCOVERY_LOCATION,
+    DISCOVERY_NAME,
+    DISCOVERY_ST,
+    DISCOVERY_UDN,
+    DISCOVERY_USN,
+    DOMAIN,
+    LOGGER as _LOGGER,
+)
 from .device import Device
 
-config_entry_flow.register_discovery_flow(
-    DOMAIN, "UPnP/IGD", Device.async_discover, config_entries.CONN_CLASS_LOCAL_POLL
-)
+
+class UpnpFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
+    """Handle a UPnP/IGD config flow."""
+
+    VERSION = 1
+    CONNECTION_CLASS = config_entries.CONN_CLASS_LOCAL_POLL
+
+    # Paths:
+    # - ssdp(discovery_info) --> ssdp_confirm(None) --> ssdp_confirm({}) --> create_entry()
+    # - user(None): scan --> user({...}) --> create_entry()
+    # - import(None) --> create_entry()
+
+    def __init__(self):
+        """Initialize the UPnP/IGD config flow."""
+        self._discoveries: Mapping = None
+
+    async def async_step_user(self, user_input: Optional[Mapping] = None):
+        """Handle a flow start."""
+        _LOGGER.debug("async_step_user: user_input: %s", user_input)
+        # This uses DISCOVERY_USN as the identifier for the device.
+
+        if user_input is not None:
+            # Ensure wanted device was discovered.
+            matching_discoveries = [
+                discovery
+                for discovery in self._discoveries
+                if discovery[DISCOVERY_USN] == user_input["usn"]
+            ]
+            if not matching_discoveries:
+                return self.async_abort(reason="no_devices_discovered")
+
+            discovery = matching_discoveries[0]
+            await self.async_set_unique_id(
+                discovery[DISCOVERY_USN], raise_on_progress=False
+            )
+            return await self._async_create_entry_from_data(discovery)
+
+        # Discover devices.
+        discoveries = await Device.async_discover(self.hass)
+
+        # Store discoveries which have not been configured, add name for each discovery.
+        current_usns = {entry.unique_id for entry in self._async_current_entries()}
+        self._discoveries = [
+            {
+                **discovery,
+                DISCOVERY_NAME: await self._async_get_name_for_discovery(discovery),
+            }
+            for discovery in discoveries
+            if discovery[DISCOVERY_USN] not in current_usns
+        ]
+
+        # Ensure anything to add.
+        if not self._discoveries:
+            return self.async_abort(reason="no_devices_found")
+
+        data_schema = vol.Schema(
+            {
+                vol.Required("usn"): vol.In(
+                    {
+                        discovery[DISCOVERY_USN]: discovery[DISCOVERY_NAME]
+                        for discovery in self._discoveries
+                    }
+                ),
+            }
+        )
+        return self.async_show_form(step_id="user", data_schema=data_schema,)
+
+    async def async_step_import(self, import_info: Optional[Mapping]):
+        """Import a new UPnP/IGD device as a config entry.
+
+        This flow is triggered by `async_setup`. If no device has been
+        configured before, find any device and create a config_entry for it.
+        Otherwise, do nothing.
+        """
+        _LOGGER.debug("async_step_import: import_info: %s", import_info)
+
+        if import_info is None:
+            # Landed here via configuration.yaml entry.
+            # Any device already added, then abort.
+            if self._async_current_entries():
+                _LOGGER.debug("aborting, already configured")
+                return self.async_abort(reason="already_configured")
+
+        # Test if import_info isn't already configured.
+        if import_info is not None and any(
+            import_info["udn"] == entry.data[CONFIG_ENTRY_UDN]
+            and import_info["st"] == entry.data[CONFIG_ENTRY_ST]
+            for entry in self._async_current_entries()
+        ):
+            return self.async_abort(reason="already_configured")
+
+        # Discover devices.
+        self._discoveries = await Device.async_discover(self.hass)
+
+        # Ensure anything to add. If not, silently abort.
+        if not self._discoveries:
+            _LOGGER.info("No UPnP devices discovered, aborting.")
+            return self.async_abort(reason="no_devices_found")
+
+        discovery = self._discoveries[0]
+        return await self._async_create_entry_from_data(discovery)
+
+    async def async_step_ssdp(self, discovery_info: Mapping):
+        """Handle a discovered UPnP/IGD device.
+
+        This flow is triggered by the SSDP component. It will check if the
+        host is already configured and delegate to the import step if not.
+        """
+        _LOGGER.debug("async_step_ssdp: discovery_info: %s", discovery_info)
+
+        # Ensure not already configuring/configured.
+        udn = discovery_info[ssdp.ATTR_UPNP_UDN]
+        st = discovery_info[ssdp.ATTR_SSDP_ST]  # pylint: disable=invalid-name
+        usn = f"{udn}::{st}"
+        await self.async_set_unique_id(usn)
+        self._abort_if_unique_id_configured()
+
+        # Store discovery.
+        name = discovery_info.get("friendlyName", "")
+        discovery = {
+            DISCOVERY_UDN: udn,
+            DISCOVERY_ST: st,
+            DISCOVERY_NAME: name,
+        }
+        self._discoveries = [discovery]
+
+        # Ensure user recognizable.
+        # pylint: disable=no-member # https://github.com/PyCQA/pylint/issues/3167
+        self.context["title_placeholders"] = {
+            "name": name,
+        }
+
+        return await self.async_step_ssdp_confirm()
+
+    async def async_step_ssdp_confirm(self, user_input: Optional[Mapping] = None):
+        """Confirm integration via SSDP."""
+        _LOGGER.debug("async_step_ssdp_confirm: user_input: %s", user_input)
+        if user_input is None:
+            return self.async_show_form(step_id="ssdp_confirm")
+
+        discovery = self._discoveries[0]
+        return await self._async_create_entry_from_data(discovery)
+
+    async def _async_create_entry_from_data(self, discovery: Mapping):
+        """Create an entry from own _data."""
+        _LOGGER.debug("_async_create_entry_from_data: discovery: %s", discovery)
+        # Get name from device, if not found already.
+        if DISCOVERY_NAME not in discovery and DISCOVERY_LOCATION in discovery:
+            discovery[DISCOVERY_NAME] = await self._async_get_name_for_discovery(
+                discovery
+            )
+
+        title = discovery.get(DISCOVERY_NAME, "")
+        data = {
+            CONFIG_ENTRY_UDN: discovery[DISCOVERY_UDN],
+            CONFIG_ENTRY_ST: discovery[DISCOVERY_ST],
+        }
+        return self.async_create_entry(title=title, data=data)
+
+    async def _async_get_name_for_discovery(self, discovery: Mapping):
+        """Get the name of the device from a discovery."""
+        _LOGGER.debug("_async_get_name_for_discovery: discovery: %s", discovery)
+        device = await Device.async_create_device(
+            self.hass, discovery[DISCOVERY_LOCATION]
+        )
+        return device.name

homeassistant/components/upnp/const.py

@@ -20,3 +20,10 @@
 DATA_RATE_PACKETS_PER_SECOND = f"{DATA_PACKETS}/{TIME_SECONDS}"
 KIBIBYTE = 1024
 UPDATE_INTERVAL = timedelta(seconds=30)
+DISCOVERY_NAME = "name"
+DISCOVERY_LOCATION = "location"
+DISCOVERY_ST = "st"
+DISCOVERY_UDN = "udn"
+DISCOVERY_USN = "usn"
+CONFIG_ENTRY_UDN = "udn"
+CONFIG_ENTRY_ST = "st"