Add discovery support to unifiprotect (#64340)

homeassistant/components/unifiprotect/__init__.py

@@ -33,6 +33,7 @@ from .const import (
     PLATFORMS,
 )
 from .data import ProtectData
+from .discovery import async_start_discovery
 from .services import async_cleanup_services, async_setup_services
 
 _LOGGER = logging.getLogger(__name__)
@@ -43,6 +44,7 @@ SCAN_INTERVAL = timedelta(seconds=DEFAULT_SCAN_INTERVAL)
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up the UniFi Protect config entries."""
 
+    await async_start_discovery(hass)
     session = async_create_clientsession(hass, cookie_jar=CookieJar(unsafe=True))
     protect = ProtectApiClient(
         host=entry.data[CONF_HOST],

homeassistant/components/unifiprotect/config_flow.py

@@ -10,6 +10,7 @@ from pyunifiprotect.data.nvr import NVR
 import voluptuous as vol
 
 from homeassistant import config_entries
+from homeassistant.components import dhcp, ssdp
 from homeassistant.const import (
     CONF_HOST,
     CONF_ID,
@@ -21,6 +22,7 @@ from homeassistant.const import (
 from homeassistant.core import callback
 from homeassistant.data_entry_flow import FlowResult
 from homeassistant.helpers.aiohttp_client import async_create_clientsession
+from homeassistant.helpers.typing import DiscoveryInfoType
 
 from .const import (
     CONF_ALL_UPDATES,
@@ -32,6 +34,8 @@ from .const import (
     MIN_REQUIRED_PROTECT_V,
     OUTDATED_LOG_MESSAGE,
 )
+from .discovery import async_start_discovery
+from .services import _async_unifi_mac_from_hass
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -44,8 +48,77 @@ class ProtectFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
     def __init__(self) -> None:
         """Init the config flow."""
         super().__init__()
-
         self.entry: config_entries.ConfigEntry | None = None
+        self._discovered_device: dict[str, str] = {}
+
+    async def async_step_dhcp(self, discovery_info: dhcp.DhcpServiceInfo) -> FlowResult:
+        """Handle discovery via dhcp."""
+        _LOGGER.debug("Starting discovery via: %s", discovery_info)
+        return await self._async_discovery_handoff()
+
+    async def async_step_ssdp(self, discovery_info: ssdp.SsdpServiceInfo) -> FlowResult:
+        """Handle a discovered UniFi device."""
+        _LOGGER.debug("Starting discovery via: %s", discovery_info)
+        return await self._async_discovery_handoff()
+
+    async def _async_discovery_handoff(self) -> FlowResult:
+        """Ensure discovery is active."""
+        # Discovery requires an additional check so we use
+        # SSDP and DHCP to tell us to start it so it only
+        # runs on networks where unifi devices are present.
+        await async_start_discovery(self.hass)
+        return self.async_abort(reason="discovery_started")
+
+    async def async_step_discovery(
+        self, discovery_info: DiscoveryInfoType
+    ) -> FlowResult:
+        """Handle discovery."""
+        self._discovered_device = discovery_info
+        mac = _async_unifi_mac_from_hass(discovery_info["mac"])
+        await self.async_set_unique_id(mac)
+        self._abort_if_unique_id_configured(
+            updates={CONF_HOST: discovery_info["ip_address"]}
+        )
+        return await self.async_step_discovery_confirm()
+
+    async def async_step_discovery_confirm(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Confirm discovery."""
+        errors: dict[str, str] = {}
+        discovery_info = self._discovered_device
+        if user_input is not None:
+            user_input[CONF_HOST] = discovery_info["ip_address"]
+            user_input[CONF_PORT] = DEFAULT_PORT
+            nvr_data, errors = await self._async_get_nvr_data(user_input)
+            if nvr_data and not errors:
+                return self._async_create_entry(nvr_data.name, user_input)
+
+        placeholders = {
+            "name": discovery_info["hostname"]
+            or discovery_info["platform"]
+            or f"NVR {discovery_info['mac']}",
+            "ip_address": discovery_info["ip_address"],
+        }
+        self.context["title_placeholders"] = placeholders
+        user_input = user_input or {}
+        return self.async_show_form(
+            step_id="discovery_confirm",
+            description_placeholders=placeholders,
+            data_schema=vol.Schema(
+                {
+                    vol.Required(
+                        CONF_VERIFY_SSL,
+                        default=user_input.get(CONF_VERIFY_SSL, DEFAULT_VERIFY_SSL),
+                    ): bool,
+                    vol.Required(
+                        CONF_USERNAME, default=user_input.get(CONF_USERNAME)
+                    ): str,
+                    vol.Required(CONF_PASSWORD): str,
+                }
+            ),
+            errors=errors,
+        )
 
     @staticmethod
     @callback

homeassistant/components/unifiprotect/discovery.py

@@ -0,0 +1,66 @@
+"""The unifiprotect integration discovery."""
+from __future__ import annotations
+
+import asyncio
+from datetime import timedelta
+import logging
+from typing import Any
+
+from unifi_discovery import AIOUnifiScanner, UnifiDevice, UnifiService
+
+from homeassistant import config_entries
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.event import async_track_time_interval
+
+from .const import DOMAIN
+
+_LOGGER = logging.getLogger(__name__)
+
+DISCOVERY = "discovery"
+DISCOVERY_INTERVAL = timedelta(minutes=60)
+
+
+async def async_start_discovery(hass: HomeAssistant) -> None:
+    """Start discovery."""
+    domain_data = hass.data.setdefault(DOMAIN, {})
+    if DISCOVERY in domain_data:
+        return
+    domain_data[DISCOVERY] = True
+
+    async def _async_discovery(*_: Any) -> None:
+        async_trigger_discovery(hass, await async_discover_devices(hass))
+
+    # Do not block startup since discovery takes 31s or more
+    asyncio.create_task(_async_discovery())
+
+    async_track_time_interval(hass, _async_discovery, DISCOVERY_INTERVAL)
+
+
+async def async_discover_devices(hass: HomeAssistant) -> list[UnifiDevice]:
+    """Discover devices."""
+    scanner = AIOUnifiScanner()
+    devices = await scanner.async_scan()
+    _LOGGER.debug("Found devices: %s", devices)
+    return devices
+
+
+@callback
+def async_trigger_discovery(
+    hass: HomeAssistant,
+    discovered_devices: list[UnifiDevice],
+) -> None:
+    """Trigger config flows for discovered devices."""
+    for device in discovered_devices:
+        if device.services[UnifiService.Protect] and device.hw_addr:
+            hass.async_create_task(
+                hass.config_entries.flow.async_init(
+                    DOMAIN,
+                    context={"source": config_entries.SOURCE_DISCOVERY},
+                    data={
+                        "ip_address": device.source_ip,
+                        "mac": device.hw_addr,
+                        "hostname": device.hostname,  # can be None
+                        "platform": device.platform,  # can be None
+                    },
+                )
+            )

homeassistant/components/unifiprotect/manifest.json

@@ -4,7 +4,7 @@
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/unifiprotect",
   "requirements": [
-    "pyunifiprotect==3.1.1"
+    "pyunifiprotect==3.1.1", "unifi-discovery==1.0.0"
   ],
   "dependencies": [
     "http"
@@ -15,5 +15,48 @@
     "@bdraco"
   ],
   "quality_scale": "platinum",
-  "iot_class": "local_push"
+  "iot_class": "local_push",
+  "dhcp": [
+    {
+      "macaddress": "B4FBE4*"
+    },
+    {
+      "macaddress": "802AA8*"
+    },
+    {
+      "macaddress": "F09FC2*"
+    },
+    {
+      "macaddress": "68D79A*"
+    },
+    {
+      "macaddress": "18E829*"
+    },
+    {
+      "macaddress": "245A4C*"
+    },
+    {
+      "macaddress": "784558*"
+    },
+    {
+      "macaddress": "E063DA*"
+    },
+    {
+      "macaddress": "265A4C*"
+    }
+  ],
+  "ssdp": [
+    {
+      "manufacturer": "Ubiquiti Networks",
+      "modelDescription": "UniFi Dream Machine"
+    },
+    {
+      "manufacturer": "Ubiquiti Networks",
+      "modelDescription": "UniFi Dream Machine Pro"
+    },
+    {
+      "manufacturer": "Ubiquiti Networks",
+      "modelDescription": "UniFi Dream Machine SE"
+    }    
+  ]
 }

homeassistant/components/unifiprotect/strings.json

@@ -1,5 +1,6 @@
 {
   "config": {
+    "flow_title": "{name} ({ip_address})",
     "step": {
       "user": {
         "title": "UniFi Protect Setup",
@@ -19,6 +20,14 @@
           "username": "[%key:common::config_flow::data::username%]",
           "password": "[%key:common::config_flow::data::password%]"
         }
+      },
+      "discovery_confirm": {
+        "description": "Do you want to setup {name} ({ip_address})?",
+        "data": {
+          "verify_ssl": "[%key:common::config_flow::data::verify_ssl%]",
+          "username": "[%key:common::config_flow::data::username%]",
+          "password": "[%key:common::config_flow::data::password%]"        
+        }
       }      
     },
     "error": {
@@ -27,7 +36,8 @@
       "protect_version": "Minimum required version is v1.20.0. Please upgrade UniFi Protect and then retry."
     },
     "abort": {
-      "already_configured": "[%key:common::config_flow::abort::already_configured_device%]"
+      "already_configured": "[%key:common::config_flow::abort::already_configured_device%]",
+      "discovery_started": "Discovery started"
     }
   },
   "options": {

homeassistant/components/unifiprotect/translations/en.json

@@ -1,15 +1,24 @@
 {
     "config": {
         "abort": {
-            "already_configured": "Device is already configured"
+            "already_configured": "Device is already configured",
+            "discovery_started": "Discovery started"
         },
         "error": {
             "cannot_connect": "Failed to connect",
             "invalid_auth": "Invalid authentication",
-            "protect_version": "Minimum required version is v1.20.0. Please upgrade UniFi Protect and then retry.",
+            "protect_version": "Minimum required version is v1.20.0. Please upgrade UniFi Protect and then retry."
-            "unknown": "Unexpected error"
         },
+        "flow_title": "{name} ({ip_address})",
         "step": {
+            "discovery_confirm": {
+                "data": {
+                    "password": "Password",
+                    "username": "Username",
+                    "verify_ssl": "Verify SSL certificate"
+                },
+                "description": "Do you want to setup {name} ({ip_address})?"
+            },
             "reauth_confirm": {
                 "data": {
                     "host": "IP/Host of UniFi Protect Server",

homeassistant/generated/dhcp.py

@@ -546,6 +546,42 @@ DHCP = [
         "domain": "twinkly",
         "hostname": "twinkly_*"
     },
+    {
+        "domain": "unifiprotect",
+        "macaddress": "B4FBE4*"
+    },
+    {
+        "domain": "unifiprotect",
+        "macaddress": "802AA8*"
+    },
+    {
+        "domain": "unifiprotect",
+        "macaddress": "F09FC2*"
+    },
+    {
+        "domain": "unifiprotect",
+        "macaddress": "68D79A*"
+    },
+    {
+        "domain": "unifiprotect",
+        "macaddress": "18E829*"
+    },
+    {
+        "domain": "unifiprotect",
+        "macaddress": "245A4C*"
+    },
+    {
+        "domain": "unifiprotect",
+        "macaddress": "784558*"
+    },
+    {
+        "domain": "unifiprotect",
+        "macaddress": "E063DA*"
+    },
+    {
+        "domain": "unifiprotect",
+        "macaddress": "265A4C*"
+    },
     {
         "domain": "verisure",
         "macaddress": "0023C1*"

homeassistant/generated/ssdp.py

@@ -238,6 +238,20 @@ SSDP = {
             "modelDescription": "UniFi Dream Machine Pro"
         }
     ],
+    "unifiprotect": [
+        {
+            "manufacturer": "Ubiquiti Networks",
+            "modelDescription": "UniFi Dream Machine"
+        },
+        {
+            "manufacturer": "Ubiquiti Networks",
+            "modelDescription": "UniFi Dream Machine Pro"
+        },
+        {
+            "manufacturer": "Ubiquiti Networks",
+            "modelDescription": "UniFi Dream Machine SE"
+        }
+    ],
     "upnp": [
         {
             "st": "urn:schemas-upnp-org:device:InternetGatewayDevice:1"

requirements_all.txt

@@ -2391,6 +2391,9 @@ twilio==6.32.0
 # homeassistant.components.rainforest_eagle
 uEagle==0.0.2
 
+# homeassistant.components.unifiprotect
+unifi-discovery==1.0.0
+
 # homeassistant.components.unifiled
 unifiled==0.11
 

requirements_test_all.txt

@@ -1449,6 +1449,9 @@ twilio==6.32.0
 # homeassistant.components.rainforest_eagle
 uEagle==0.0.2
 
+# homeassistant.components.unifiprotect
+unifi-discovery==1.0.0
+
 # homeassistant.components.upb
 upb_lib==0.4.12
 

tests/components/unifiprotect/__init__.py

@@ -1 +1,36 @@
 """Tests for the UniFi Protect integration."""
+
+from contextlib import contextmanager
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from unifi_discovery import AIOUnifiScanner, UnifiDevice, UnifiService
+
+DEVICE_HOSTNAME = "unvr"
+DEVICE_IP_ADDRESS = "127.0.0.1"
+DEVICE_MAC_ADDRESS = "aa:bb:cc:dd:ee:ff"
+
+
+UNIFI_DISCOVERY = UnifiDevice(
+    source_ip=DEVICE_IP_ADDRESS,
+    hw_addr=DEVICE_MAC_ADDRESS,
+    platform=DEVICE_HOSTNAME,
+    hostname=DEVICE_HOSTNAME,
+    services={UnifiService.Protect: True},
+)
+
+
+def _patch_discovery(device=None, no_device=False):
+    mock_aio_discovery = MagicMock(auto_spec=AIOUnifiScanner)
+    scanner_return = [] if no_device else [device or UNIFI_DISCOVERY]
+    mock_aio_discovery.async_scan = AsyncMock(return_value=scanner_return)
+    mock_aio_discovery.found_devices = scanner_return
+
+    @contextmanager
+    def _patcher():
+        with patch(
+            "homeassistant.components.unifiprotect.discovery.AIOUnifiScanner",
+            return_value=mock_aio_discovery,
+        ):
+            yield
+
+    return _patcher()

tests/components/unifiprotect/conftest.py

@@ -23,6 +23,8 @@ from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity import EntityDescription
 import homeassistant.util.dt as dt_util
 
+from . import _patch_discovery
+
 from tests.common import MockConfigEntry, async_fire_time_changed, load_fixture
 
 MAC_ADDR = "aa:bb:cc:dd:ee:ff"
@@ -73,6 +75,24 @@ def mock_nvr_fixture():
     NVR.__config__.validate_assignment = True
 
 
+@pytest.fixture(name="mock_ufp_config_entry")
+def mock_ufp_config_entry():
+    """Mock the unifiprotect config entry."""
+
+    return MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            "host": "1.1.1.1",
+            "username": "test-username",
+            "password": "test-password",
+            "id": "UnifiProtect",
+            "port": 443,
+            "verify_ssl": False,
+        },
+        version=2,
+    )
+
+
 @pytest.fixture(name="mock_old_nvr")
 def mock_old_nvr_fixture():
     """Mock UniFi Protect Camera device."""
@@ -122,28 +142,20 @@ def mock_client(mock_bootstrap: MockBootstrap):
 
 @pytest.fixture
 def mock_entry(
-    hass: HomeAssistant, mock_client  # pylint: disable=redefined-outer-name
+    hass: HomeAssistant,
+    mock_ufp_config_entry: MockConfigEntry,
+    mock_client,  # pylint: disable=redefined-outer-name
 ):
     """Mock ProtectApiClient for testing."""
 
-    with patch("homeassistant.components.unifiprotect.ProtectApiClient") as mock_api:
+    with _patch_discovery(no_device=True), patch(
-        mock_config = MockConfigEntry(
+        "homeassistant.components.unifiprotect.ProtectApiClient"
-            domain=DOMAIN,
+    ) as mock_api:
-            data={
+        mock_ufp_config_entry.add_to_hass(hass)
-                "host": "1.1.1.1",
-                "username": "test-username",
-                "password": "test-password",
-                "id": "UnifiProtect",
-                "port": 443,
-                "verify_ssl": False,
-            },
-            version=2,
-        )
-        mock_config.add_to_hass(hass)
 
         mock_api.return_value = mock_client
 
-        yield MockEntityFixture(mock_config, mock_client)
+        yield MockEntityFixture(mock_ufp_config_entry, mock_client)
 
 
 @pytest.fixture

tests/components/unifiprotect/test_config_flow.py

@@ -3,10 +3,12 @@ from __future__ import annotations
 
 from unittest.mock import patch
 
+import pytest
 from pyunifiprotect import NotAuthorized, NvrError
 from pyunifiprotect.data.nvr import NVR
 
 from homeassistant import config_entries
+from homeassistant.components import dhcp, ssdp
 from homeassistant.components.unifiprotect.const import (
     CONF_ALL_UPDATES,
     CONF_DISABLE_RTSP,
@@ -21,10 +23,35 @@ from homeassistant.data_entry_flow import (
 )
 from homeassistant.helpers import device_registry as dr
 
+from . import DEVICE_HOSTNAME, DEVICE_IP_ADDRESS, DEVICE_MAC_ADDRESS, _patch_discovery
 from .conftest import MAC_ADDR
 
 from tests.common import MockConfigEntry
 
+DHCP_DISCOVERY = dhcp.DhcpServiceInfo(
+    hostname=DEVICE_HOSTNAME,
+    ip=DEVICE_IP_ADDRESS,
+    macaddress=DEVICE_MAC_ADDRESS,
+)
+SSDP_DISCOVERY = (
+    ssdp.SsdpServiceInfo(
+        ssdp_usn="mock_usn",
+        ssdp_st="mock_st",
+        ssdp_location=f"http://{DEVICE_IP_ADDRESS}:41417/rootDesc.xml",
+        upnp={
+            "friendlyName": "UniFi Dream Machine",
+            "modelDescription": "UniFi Dream Machine Pro",
+            "serialNumber": DEVICE_MAC_ADDRESS,
+        },
+    ),
+)
+UNIFI_DISCOVERY_DICT = {
+    "ip_address": DEVICE_IP_ADDRESS,
+    "mac": DEVICE_MAC_ADDRESS,
+    "hostname": DEVICE_HOSTNAME,
+    "platform": DEVICE_HOSTNAME,
+}
+
 
 async def test_form(hass: HomeAssistant, mock_nvr: NVR) -> None:
     """Test we get the form."""
@@ -208,7 +235,9 @@ async def test_form_options(hass: HomeAssistant, mock_client) -> None:
     )
     mock_config.add_to_hass(hass)
 
-    with patch("homeassistant.components.unifiprotect.ProtectApiClient") as mock_api:
+    with _patch_discovery(), patch(
+        "homeassistant.components.unifiprotect.ProtectApiClient"
+    ) as mock_api:
         mock_api.return_value = mock_client
 
         await hass.config_entries.async_setup(mock_config.entry_id)
@@ -231,3 +260,73 @@ async def test_form_options(hass: HomeAssistant, mock_client) -> None:
         "disable_rtsp": True,
         "override_connection_host": True,
     }
+
+
+@pytest.mark.parametrize(
+    "source, data",
+    [
+        (config_entries.SOURCE_DHCP, DHCP_DISCOVERY),
+        (config_entries.SOURCE_SSDP, SSDP_DISCOVERY),
+    ],
+)
+async def test_discovered_by_ssdp_or_dhcp(
+    hass: HomeAssistant, source: str, data: dhcp.DhcpServiceInfo | ssdp.SsdpServiceInfo
+) -> None:
+    """Test we handoff to unifi-discovery when discovered via ssdp or dhcp."""
+
+    with _patch_discovery():
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": source},
+            data=data,
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == RESULT_TYPE_ABORT
+    assert result["reason"] == "discovery_started"
+
+
+async def test_discovered_by_unifi_discovery(
+    hass: HomeAssistant, mock_nvr: NVR
+) -> None:
+    """Test a discovery from unifi-discovery."""
+
+    with _patch_discovery():
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={"source": config_entries.SOURCE_DISCOVERY},
+            data=UNIFI_DISCOVERY_DICT,
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] == RESULT_TYPE_FORM
+    assert result["step_id"] == "discovery_confirm"
+    assert not result["errors"]
+
+    with patch(
+        "homeassistant.components.unifiprotect.config_flow.ProtectApiClient.get_nvr",
+        return_value=mock_nvr,
+    ), patch(
+        "homeassistant.components.unifiprotect.async_setup_entry",
+        return_value=True,
+    ) as mock_setup_entry:
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                "username": "test-username",
+                "password": "test-password",
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result2["type"] == RESULT_TYPE_CREATE_ENTRY
+    assert result2["title"] == "UnifiProtect"
+    assert result2["data"] == {
+        "host": DEVICE_IP_ADDRESS,
+        "username": "test-username",
+        "password": "test-password",
+        "id": "UnifiProtect",
+        "port": 443,
+        "verify_ssl": False,
+    }
+    assert len(mock_setup_entry.mock_calls) == 1

tests/components/unifiprotect/test_init.py

@@ -7,9 +7,10 @@ from pyunifiprotect import NotAuthorized, NvrError
 from pyunifiprotect.data import NVR
 
 from homeassistant.components.unifiprotect.const import CONF_DISABLE_RTSP, DOMAIN
-from homeassistant.config_entries import ConfigEntryState
+from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.core import HomeAssistant
 
+from . import _patch_discovery
 from .conftest import MockBootstrap, MockEntityFixture
 
 from tests.common import MockConfigEntry
@@ -156,3 +157,21 @@ async def test_setup_failed_auth(hass: HomeAssistant, mock_entry: MockEntityFixt
     await hass.config_entries.async_setup(mock_entry.entry.entry_id)
     assert mock_entry.entry.state == ConfigEntryState.SETUP_ERROR
     assert not mock_entry.api.update.called
+
+
+async def test_setup_starts_discovery(
+    hass: HomeAssistant, mock_ufp_config_entry: ConfigEntry, mock_client
+):
+    """Test setting up will start discovery."""
+    with _patch_discovery(), patch(
+        "homeassistant.components.unifiprotect.ProtectApiClient"
+    ) as mock_api:
+        mock_ufp_config_entry.add_to_hass(hass)
+        mock_api.return_value = mock_client
+        mock_entry = MockEntityFixture(mock_ufp_config_entry, mock_client)
+
+        await hass.config_entries.async_setup(mock_entry.entry.entry_id)
+        await hass.async_block_till_done()
+        assert mock_entry.entry.state == ConfigEntryState.LOADED
+        await hass.async_block_till_done()
+        assert len(hass.config_entries.flow.async_progress_by_handler(DOMAIN)) == 1