Add network and callback support to SSDP (#51019)

Co-authored-by: Ruslan Sayfutdinov <ruslan@sayfutdinov.com>
Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>

homeassistant/components/ssdp/__init__.py

@@ -4,18 +4,23 @@ from __future__ import annotations
 import asyncio
 from collections.abc import Mapping
 from datetime import timedelta
+from ipaddress import IPv4Address, IPv6Address
 import logging
-from typing import Any
+from typing import Any, Callable
 
-import aiohttp
-from async_upnp_client.search import async_search
-from defusedxml import ElementTree
-from netdisco import ssdp, util
+from async_upnp_client.search import SSDPListener
+from async_upnp_client.utils import CaseInsensitiveDict
 
+from homeassistant import config_entries
+from homeassistant.components import network
 from homeassistant.const import EVENT_HOMEASSISTANT_STARTED, EVENT_HOMEASSISTANT_STOP
-from homeassistant.core import callback
+from homeassistant.core import CoreState, HomeAssistant, callback as core_callback
 from homeassistant.helpers.event import async_track_time_interval
-from homeassistant.loader import async_get_ssdp
+from homeassistant.helpers.typing import ConfigType
+from homeassistant.loader import async_get_ssdp, bind_hass
+
+from .descriptions import DescriptionManager
+from .flow import FlowDispatcher, SSDPFlow
 
 DOMAIN = "ssdp"
 SCAN_INTERVAL = timedelta(seconds=60)
@@ -40,188 +45,321 @@ ATTR_UPNP_UDN = "UDN"
 ATTR_UPNP_UPC = "UPC"
 ATTR_UPNP_PRESENTATION_URL = "presentationURL"
 
+
+DISCOVERY_MAPPING = {
+    "usn": ATTR_SSDP_USN,
+    "ext": ATTR_SSDP_EXT,
+    "server": ATTR_SSDP_SERVER,
+    "st": ATTR_SSDP_ST,
+    "location": ATTR_SSDP_LOCATION,
+}
+
+
 _LOGGER = logging.getLogger(__name__)
 
 
-async def async_setup(hass, config):
+@bind_hass
+def async_register_callback(
+    hass: HomeAssistant,
+    callback: Callable[[dict], None],
+    match_dict: None | dict[str, str] = None,
+) -> Callable[[], None]:
+    """Register to receive a callback on ssdp broadcast.
+
+    Returns a callback that can be used to cancel the registration.
+    """
+    scanner: Scanner = hass.data[DOMAIN]
+    return scanner.async_register_callback(callback, match_dict)
+
+
+@bind_hass
+def async_get_discovery_info_by_udn_st(  # pylint: disable=invalid-name
+    hass: HomeAssistant, udn: str, st: str
+) -> dict[str, str] | None:
+    """Fetch the discovery info cache."""
+    scanner: Scanner = hass.data[DOMAIN]
+    return scanner.async_get_discovery_info_by_udn_st(udn, st)
+
+
+@bind_hass
+def async_get_discovery_info_by_st(  # pylint: disable=invalid-name
+    hass: HomeAssistant, st: str
+) -> list[dict[str, str]]:
+    """Fetch all the entries matching the st."""
+    scanner: Scanner = hass.data[DOMAIN]
+    return scanner.async_get_discovery_info_by_st(st)
+
+
+@bind_hass
+def async_get_discovery_info_by_udn(
+    hass: HomeAssistant, udn: str
+) -> list[dict[str, str]]:
+    """Fetch all the entries matching the udn."""
+    scanner: Scanner = hass.data[DOMAIN]
+    return scanner.async_get_discovery_info_by_udn(udn)
+
+
+async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up the SSDP integration."""
 
-    async def _async_initialize(_):
-        scanner = Scanner(hass, await async_get_ssdp(hass))
-        await scanner.async_scan(None)
-        cancel_scan = async_track_time_interval(hass, scanner.async_scan, SCAN_INTERVAL)
+    scanner = hass.data[DOMAIN] = Scanner(hass, await async_get_ssdp(hass))
 
-        @callback
-        def _async_stop_scans(event):
-            cancel_scan()
-
-        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _async_stop_scans)
-
-    hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STARTED, _async_initialize)
+    asyncio.create_task(scanner.async_start())
 
     return True
 
 
+@core_callback
+def _async_use_default_interface(adapters: list[network.Adapter]) -> bool:
+    for adapter in adapters:
+        if adapter["enabled"] and not adapter["default"]:
+            return False
+    return True
+
+
+@core_callback
+def _async_process_callbacks(
+    callbacks: list[Callable[[dict], None]], discovery_info: dict[str, str]
+) -> None:
+    for callback in callbacks:
+        try:
+            callback(discovery_info)
+        except Exception:  # pylint: disable=broad-except
+            _LOGGER.exception("Failed to callback info: %s", discovery_info)
+
+
 class Scanner:
     """Class to manage SSDP scanning."""
 
-    def __init__(self, hass, integration_matchers):
+    def __init__(
+        self, hass: HomeAssistant, integration_matchers: dict[str, list[dict[str, str]]]
+    ) -> None:
         """Initialize class."""
         self.hass = hass
-        self.seen = set()
-        self._entries = []
+        self.seen: set[tuple[str, str]] = set()
+        self.cache: dict[tuple[str, str], Mapping[str, str]] = {}
         self._integration_matchers = integration_matchers
-        self._description_cache = {}
+        self._cancel_scan: Callable[[], None] | None = None
+        self._ssdp_listeners: list[SSDPListener] = []
+        self._callbacks: list[tuple[Callable[[dict], None], dict[str, str]]] = []
+        self.flow_dispatcher: FlowDispatcher | None = None
+        self.description_manager: DescriptionManager | None = None
 
-    async def _on_ssdp_response(self, data: Mapping[str, Any]) -> None:
-        """Process an ssdp response."""
-        self.async_store_entry(
-            ssdp.UPNPEntry({key.lower(): item for key, item in data.items()})
+    @core_callback
+    def async_register_callback(
+        self, callback: Callable[[dict], None], match_dict: None | dict[str, str] = None
+    ) -> Callable[[], None]:
+        """Register a callback."""
+        if match_dict is None:
+            match_dict = {}
+
+        # Make sure any entries that happened
+        # before the callback was registered are fired
+        if self.hass.state != CoreState.running:
+            for headers in self.cache.values():
+                self._async_callback_if_match(callback, headers, match_dict)
+
+        callback_entry = (callback, match_dict)
+        self._callbacks.append(callback_entry)
+
+        @core_callback
+        def _async_remove_callback() -> None:
+            self._callbacks.remove(callback_entry)
+
+        return _async_remove_callback
+
+    @core_callback
+    def _async_callback_if_match(
+        self,
+        callback: Callable[[dict], None],
+        headers: Mapping[str, str],
+        match_dict: dict[str, str],
+    ) -> None:
+        """Fire a callback if info matches the match dict."""
+        if not all(headers.get(k) == v for (k, v) in match_dict.items()):
+            return
+        _async_process_callbacks(
+            [callback], self._async_headers_to_discovery_info(headers)
         )
 
-    @callback
-    def async_store_entry(self, entry):
-        """Save an entry for later processing."""
-        self._entries.append(entry)
+    @core_callback
+    def async_stop(self, *_: Any) -> None:
+        """Stop the scanner."""
+        assert self._cancel_scan is not None
+        self._cancel_scan()
+        for listener in self._ssdp_listeners:
+            listener.async_stop()
+        self._ssdp_listeners = []
 
-    async def async_scan(self, _):
+    async def _async_build_source_set(self) -> set[IPv4Address | IPv6Address]:
+        """Build the list of ssdp sources."""
+        adapters = await network.async_get_adapters(self.hass)
+        sources: set[IPv4Address | IPv6Address] = set()
+        if _async_use_default_interface(adapters):
+            sources.add(IPv4Address("0.0.0.0"))
+            return sources
+
+        for adapter in adapters:
+            if not adapter["enabled"]:
+                continue
+            if adapter["ipv4"]:
+                ipv4 = adapter["ipv4"][0]
+                sources.add(IPv4Address(ipv4["address"]))
+            if adapter["ipv6"]:
+                ipv6 = adapter["ipv6"][0]
+                # With python 3.9 add scope_ids can be
+                # added by enumerating adapter["ipv6"]s
+                # IPv6Address(f"::%{ipv6['scope_id']}")
+                sources.add(IPv6Address(ipv6["address"]))
+
+        return sources
+
+    @core_callback
+    def async_scan(self, *_: Any) -> None:
         """Scan for new entries."""
+        for listener in self._ssdp_listeners:
+            listener.async_search()
 
-        await async_search(async_callback=self._on_ssdp_response)
-        await self._process_entries()
-
-        # We clear the cache after each run. We track discovered entries
-        # so will never need a description twice.
-        self._description_cache.clear()
-        self._entries.clear()
-
-    async def _process_entries(self):
-        """Process SSDP entries."""
-        entries_to_process = []
-        unseen_locations = set()
-
-        for entry in self._entries:
-            key = (entry.st, entry.location)
-
-            if key in self.seen:
-                continue
-
-            self.seen.add(key)
-
-            entries_to_process.append(entry)
-
-            if (
-                entry.location is not None
-                and entry.location not in self._description_cache
-            ):
-                unseen_locations.add(entry.location)
-
-        if not entries_to_process:
-            return
-
-        if unseen_locations:
-            await self._fetch_descriptions(list(unseen_locations))
-
-        tasks = []
-
-        for entry in entries_to_process:
-            info, domains = self._process_entry(entry)
-            for domain in domains:
-                _LOGGER.debug("Discovered %s at %s", domain, entry.location)
-                tasks.append(
-                    self.hass.config_entries.flow.async_init(
-                        domain, context={"source": DOMAIN}, data=info
-                    )
+    async def async_start(self) -> None:
+        """Start the scanner."""
+        self.description_manager = DescriptionManager(self.hass)
+        self.flow_dispatcher = FlowDispatcher(self.hass)
+        for source_ip in await self._async_build_source_set():
+            self._ssdp_listeners.append(
+                SSDPListener(
+                    async_callback=self._async_process_entry, source_ip=source_ip
                 )
-
-        if tasks:
-            await asyncio.gather(*tasks)
-
-    async def _fetch_descriptions(self, locations):
-        """Fetch descriptions from locations."""
-
-        for idx, result in enumerate(
-            await asyncio.gather(
-                *[self._fetch_description(location) for location in locations],
-                return_exceptions=True,
             )
-        ):
-            location = locations[idx]
 
-            if isinstance(result, Exception):
-                _LOGGER.exception(
-                    "Failed to fetch ssdp data from: %s", location, exc_info=result
-                )
-                continue
+        self.hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, self.async_stop)
+        self.hass.bus.async_listen_once(
+            EVENT_HOMEASSISTANT_STARTED, self.flow_dispatcher.async_start
+        )
+        await asyncio.gather(
+            *[listener.async_start() for listener in self._ssdp_listeners]
+        )
+        self._cancel_scan = async_track_time_interval(
+            self.hass, self.async_scan, SCAN_INTERVAL
+        )
 
-            self._description_cache[location] = result
-
-    def _process_entry(self, entry):
-        """Process a single entry."""
-
-        info = {"st": entry.st}
-        for key in "usn", "ext", "server":
-            if key in entry.values:
-                info[key] = entry.values[key]
-
-        if entry.location:
-            # Multiple entries usually share same location. Make sure
-            # we fetch it only once.
-            info_req = self._description_cache.get(entry.location)
-            if info_req is None:
-                return (None, [])
-
-            info.update(info_req)
+    @core_callback
+    def _async_get_matching_callbacks(
+        self, headers: Mapping[str, str]
+    ) -> list[Callable[[dict], None]]:
+        """Return a list of callbacks that match."""
+        return [
+            callback
+            for callback, match_dict in self._callbacks
+            if all(headers.get(k) == v for (k, v) in match_dict.items())
+        ]
 
+    @core_callback
+    def _async_matching_domains(self, info_with_req: CaseInsensitiveDict) -> set[str]:
         domains = set()
         for domain, matchers in self._integration_matchers.items():
             for matcher in matchers:
-                if all(info.get(k) == v for (k, v) in matcher.items()):
+                if all(info_with_req.get(k) == v for (k, v) in matcher.items()):
                     domains.add(domain)
+        return domains
 
-        if domains:
-            return (info_from_entry(entry, info), domains)
+    async def _async_process_entry(self, headers: Mapping[str, str]) -> None:
+        """Process SSDP entries."""
+        _LOGGER.debug("_async_process_entry: %s", headers)
+        if "st" not in headers or "location" not in headers:
+            return
+        h_st = headers["st"]
+        h_location = headers["location"]
+        key = (h_st, h_location)
 
-        return (None, [])
+        if udn := _udn_from_usn(headers.get("usn")):
+            self.cache[(udn, h_st)] = headers
 
-    async def _fetch_description(self, xml_location):
-        """Fetch an XML description."""
-        session = self.hass.helpers.aiohttp_client.async_get_clientsession()
-        try:
-            for _ in range(2):
-                resp = await session.get(xml_location, timeout=5)
-                xml = await resp.text(errors="replace")
-                # Samsung Smart TV sometimes returns an empty document the
-                # first time. Retry once.
-                if xml:
-                    break
-        except (aiohttp.ClientError, asyncio.TimeoutError) as err:
-            _LOGGER.debug("Error fetching %s: %s", xml_location, err)
-            return {}
+        callbacks = self._async_get_matching_callbacks(headers)
+        if key in self.seen and not callbacks:
+            return
 
-        try:
-            tree = ElementTree.fromstring(xml)
-        except ElementTree.ParseError as err:
-            _LOGGER.debug("Error parsing %s: %s", xml_location, err)
-            return {}
+        assert self.description_manager is not None
+        info_req = await self.description_manager.fetch_description(h_location) or {}
+        info_with_req = CaseInsensitiveDict(**headers, **info_req)
+        discovery_info = discovery_info_from_headers_and_request(info_with_req)
 
-        return util.etree_to_dict(tree).get("root", {}).get("device", {})
+        _async_process_callbacks(callbacks, discovery_info)
+        if key in self.seen:
+            return
+        self.seen.add(key)
+
+        for domain in self._async_matching_domains(info_with_req):
+            _LOGGER.debug("Discovered %s at %s", domain, h_location)
+            flow: SSDPFlow = {
+                "domain": domain,
+                "context": {"source": config_entries.SOURCE_SSDP},
+                "data": discovery_info,
+            }
+            assert self.flow_dispatcher is not None
+            self.flow_dispatcher.create(flow)
+
+    @core_callback
+    def _async_headers_to_discovery_info(
+        self, headers: Mapping[str, str]
+    ) -> dict[str, str]:
+        """Combine the headers and description into discovery_info.
+
+        Building this is a bit expensive so we only do it on demand.
+        """
+        assert self.description_manager is not None
+        location = headers["location"]
+        info_req = self.description_manager.async_cached_description(location) or {}
+        return discovery_info_from_headers_and_request(
+            CaseInsensitiveDict(**headers, **info_req)
+        )
+
+    @core_callback
+    def async_get_discovery_info_by_udn_st(  # pylint: disable=invalid-name
+        self, udn: str, st: str
+    ) -> dict[str, str] | None:
+        """Return discovery_info for a udn and st."""
+        if headers := self.cache.get((udn, st)):
+            return self._async_headers_to_discovery_info(headers)
+        return None
+
+    @core_callback
+    def async_get_discovery_info_by_st(  # pylint: disable=invalid-name
+        self, st: str
+    ) -> list[dict[str, str]]:
+        """Return matching discovery_infos for a st."""
+        return [
+            self._async_headers_to_discovery_info(headers)
+            for udn_st, headers in self.cache.items()
+            if udn_st[1] == st
+        ]
+
+    @core_callback
+    def async_get_discovery_info_by_udn(self, udn: str) -> list[dict[str, str]]:
+        """Return matching discovery_infos for a udn."""
+        return [
+            self._async_headers_to_discovery_info(headers)
+            for udn_st, headers in self.cache.items()
+            if udn_st[0] == udn
+        ]
 
 
-def info_from_entry(entry, device_info):
-    """Get info from an entry."""
-    info = {
-        ATTR_SSDP_LOCATION: entry.location,
-        ATTR_SSDP_ST: entry.st,
-    }
-    if device_info:
-        info.update(device_info)
-        info.pop("st", None)
-        if "usn" in info:
-            info[ATTR_SSDP_USN] = info.pop("usn")
-        if "ext" in info:
-            info[ATTR_SSDP_EXT] = info.pop("ext")
-        if "server" in info:
-            info[ATTR_SSDP_SERVER] = info.pop("server")
+def discovery_info_from_headers_and_request(
+    info_with_req: CaseInsensitiveDict,
+) -> dict[str, str]:
+    """Convert headers and description to discovery_info."""
+    info = {DISCOVERY_MAPPING.get(k.lower(), k): v for k, v in info_with_req.items()}
+
+    if ATTR_UPNP_UDN not in info and ATTR_SSDP_USN in info:
+        if udn := _udn_from_usn(info[ATTR_SSDP_USN]):
+            info[ATTR_UPNP_UDN] = udn
 
     return info
+
+
+def _udn_from_usn(usn: str | None) -> str | None:
+    """Get the UDN from the USN."""
+    if usn is None:
+        return None
+    if usn.startswith("uuid:"):
+        return usn.split("::")[0]
+    return None

homeassistant/components/ssdp/descriptions.py

@@ -0,0 +1,70 @@
+"""The SSDP integration."""
+from __future__ import annotations
+
+import asyncio
+import logging
+
+import aiohttp
+from defusedxml import ElementTree
+from netdisco import util
+
+from homeassistant.core import HomeAssistant, callback
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class DescriptionManager:
+    """Class to cache and manage fetching descriptions."""
+
+    def __init__(self, hass: HomeAssistant) -> None:
+        """Init the manager."""
+        self.hass = hass
+        self._description_cache: dict[str, None | dict[str, str]] = {}
+
+    async def fetch_description(
+        self, xml_location: str | None
+    ) -> None | dict[str, str]:
+        """Fetch the location or get it from the cache."""
+        if xml_location is None:
+            return None
+        if xml_location not in self._description_cache:
+            try:
+                self._description_cache[xml_location] = await self._fetch_description(
+                    xml_location
+                )
+            except Exception:  # pylint: disable=broad-except
+                # If it fails, cache the failure so we do not keep trying over and over
+                self._description_cache[xml_location] = None
+                _LOGGER.exception("Failed to fetch ssdp data from: %s", xml_location)
+
+        return self._description_cache[xml_location]
+
+    @callback
+    def async_cached_description(self, xml_location: str) -> None | dict[str, str]:
+        """Fetch the description from the cache."""
+        return self._description_cache[xml_location]
+
+    async def _fetch_description(self, xml_location: str) -> None | dict[str, str]:
+        """Fetch an XML description."""
+        session = self.hass.helpers.aiohttp_client.async_get_clientsession()
+        try:
+            for _ in range(2):
+                resp = await session.get(xml_location, timeout=5)
+                # Samsung Smart TV sometimes returns an empty document the
+                # first time. Retry once.
+                if xml := await resp.text(errors="replace"):
+                    break
+        except (aiohttp.ClientError, asyncio.TimeoutError) as err:
+            _LOGGER.debug("Error fetching %s: %s", xml_location, err)
+            return None
+
+        try:
+            tree = ElementTree.fromstring(xml)
+        except ElementTree.ParseError as err:
+            _LOGGER.debug("Error parsing %s: %s", xml_location, err)
+            return None
+
+        parsed: dict[str, str] = (
+            util.etree_to_dict(tree).get("root", {}).get("device", {})
+        )
+        return parsed

homeassistant/components/ssdp/flow.py

@@ -0,0 +1,50 @@
+"""The SSDP integration."""
+from __future__ import annotations
+
+from collections.abc import Coroutine
+from typing import Any, TypedDict
+
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.data_entry_flow import FlowResult
+
+
+class SSDPFlow(TypedDict):
+    """A queued ssdp discovery flow."""
+
+    domain: str
+    context: dict[str, Any]
+    data: dict
+
+
+class FlowDispatcher:
+    """Dispatch discovery flows."""
+
+    def __init__(self, hass: HomeAssistant) -> None:
+        """Init the discovery dispatcher."""
+        self.hass = hass
+        self.pending_flows: list[SSDPFlow] = []
+        self.started = False
+
+    @callback
+    def async_start(self, *_: Any) -> None:
+        """Start processing pending flows."""
+        self.started = True
+        self.hass.loop.call_soon(self._async_process_pending_flows)
+
+    def _async_process_pending_flows(self) -> None:
+        for flow in self.pending_flows:
+            self.hass.async_create_task(self._init_flow(flow))
+        self.pending_flows = []
+
+    def create(self, flow: SSDPFlow) -> None:
+        """Create and add or queue a flow."""
+        if self.started:
+            self.hass.async_create_task(self._init_flow(flow))
+        else:
+            self.pending_flows.append(flow)
+
+    def _init_flow(self, flow: SSDPFlow) -> Coroutine[None, None, FlowResult]:
+        """Create a flow."""
+        return self.hass.config_entries.flow.async_init(
+            flow["domain"], context=flow["context"], data=flow["data"]
+        )

homeassistant/components/ssdp/manifest.json

@@ -4,9 +4,9 @@
   "documentation": "https://www.home-assistant.io/integrations/ssdp",
   "requirements": [
     "defusedxml==0.7.1",
-    "netdisco==2.8.3",
     "async-upnp-client==0.18.0"
   ],
+  "dependencies": ["network"],
   "after_dependencies": ["zeroconf"],
   "codeowners": [],
   "quality_scale": "internal",

homeassistant/package_constraints.txt

@@ -21,7 +21,6 @@ home-assistant-frontend==20210528.0
 httpx==0.18.0
 ifaddr==0.1.7
 jinja2>=3.0.1
-netdisco==2.8.3
 paho-mqtt==1.5.1
 pillow==8.1.2
 pip>=8.0.3,<20.3