Prevent unifiprotect from being rediscovered on UDM-PROs (#65335)

pull/65442/head
J. Nick Koston 2022-01-31 17:27:26 -06:00 committed by Paulus Schoutsen
parent 3f8d2f3102
commit 5c3d4cb9a5
4 changed files with 256 additions and 29 deletions

View File

@ -36,7 +36,7 @@ from .const import (
OUTDATED_LOG_MESSAGE,
)
from .discovery import async_start_discovery
from .utils import _async_short_mac, _async_unifi_mac_from_hass
from .utils import _async_resolve, _async_short_mac, _async_unifi_mac_from_hass
_LOGGER = logging.getLogger(__name__)
@ -88,32 +88,35 @@ class ProtectFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
self._discovered_device = discovery_info
mac = _async_unifi_mac_from_hass(discovery_info["hw_addr"])
await self.async_set_unique_id(mac)
source_ip = discovery_info["source_ip"]
direct_connect_domain = discovery_info["direct_connect_domain"]
for entry in self._async_current_entries(include_ignore=False):
if entry.unique_id != mac:
continue
new_host = None
if (
_host_is_direct_connect(entry.data[CONF_HOST])
and discovery_info["direct_connect_domain"]
and entry.data[CONF_HOST] != discovery_info["direct_connect_domain"]
entry_host = entry.data[CONF_HOST]
entry_has_direct_connect = _host_is_direct_connect(entry_host)
if entry.unique_id == mac:
new_host = None
if (
entry_has_direct_connect
and direct_connect_domain
and entry_host != direct_connect_domain
):
new_host = direct_connect_domain
elif not entry_has_direct_connect and entry_host != source_ip:
new_host = source_ip
if new_host:
self.hass.config_entries.async_update_entry(
entry, data={**entry.data, CONF_HOST: new_host}
)
self.hass.async_create_task(
self.hass.config_entries.async_reload(entry.entry_id)
)
return self.async_abort(reason="already_configured")
if entry_host in (direct_connect_domain, source_ip) or (
entry_has_direct_connect
and (ip := await _async_resolve(self.hass, entry_host))
and ip == source_ip
):
new_host = discovery_info["direct_connect_domain"]
elif (
not _host_is_direct_connect(entry.data[CONF_HOST])
and entry.data[CONF_HOST] != discovery_info["source_ip"]
):
new_host = discovery_info["source_ip"]
if new_host:
self.hass.config_entries.async_update_entry(
entry, data={**entry.data, CONF_HOST: new_host}
)
self.hass.async_create_task(
self.hass.config_entries.async_reload(entry.entry_id)
)
return self.async_abort(reason="already_configured")
self._abort_if_unique_id_configured(
updates={CONF_HOST: discovery_info["source_ip"]}
)
return self.async_abort(reason="already_configured")
return await self.async_step_discovery_confirm()
async def async_step_discovery_confirm(

View File

@ -1,10 +1,12 @@
"""UniFi Protect Integration utils."""
from __future__ import annotations
import contextlib
from enum import Enum
import socket
from typing import Any
from homeassistant.core import callback
from homeassistant.core import HomeAssistant, callback
def get_nested_attr(obj: Any, attr: str) -> Any:
@ -33,3 +35,19 @@ def _async_unifi_mac_from_hass(mac: str) -> str:
def _async_short_mac(mac: str) -> str:
"""Get the short mac address from the full mac."""
return _async_unifi_mac_from_hass(mac)[-6:]
async def _async_resolve(hass: HomeAssistant, host: str) -> str | None:
"""Resolve a hostname to an ip."""
with contextlib.suppress(OSError):
return next(
iter(
raw[0]
for family, _, _, _, raw in await hass.loop.getaddrinfo(
host, None, type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP
)
if family == socket.AF_INET
),
None,
)
return None

View File

@ -8,6 +8,7 @@ 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"
DIRECT_CONNECT_DOMAIN = "x.ui.direct"
UNIFI_DISCOVERY = UnifiDevice(
@ -16,7 +17,7 @@ UNIFI_DISCOVERY = UnifiDevice(
platform=DEVICE_HOSTNAME,
hostname=DEVICE_HOSTNAME,
services={UnifiService.Protect: True},
direct_connect_domain="x.ui.direct",
direct_connect_domain=DIRECT_CONNECT_DOMAIN,
)

View File

@ -2,6 +2,7 @@
from __future__ import annotations
from dataclasses import asdict
import socket
from unittest.mock import patch
import pytest
@ -29,6 +30,7 @@ from . import (
DEVICE_HOSTNAME,
DEVICE_IP_ADDRESS,
DEVICE_MAC_ADDRESS,
DIRECT_CONNECT_DOMAIN,
UNIFI_DISCOVERY,
UNIFI_DISCOVERY_PARTIAL,
_patch_discovery,
@ -334,7 +336,7 @@ async def test_discovered_by_unifi_discovery_direct_connect(
assert result2["type"] == RESULT_TYPE_CREATE_ENTRY
assert result2["title"] == "UnifiProtect"
assert result2["data"] == {
"host": "x.ui.direct",
"host": DIRECT_CONNECT_DOMAIN,
"username": "test-username",
"password": "test-password",
"id": "UnifiProtect",
@ -377,7 +379,7 @@ async def test_discovered_by_unifi_discovery_direct_connect_updated(
assert result["type"] == RESULT_TYPE_ABORT
assert result["reason"] == "already_configured"
assert len(mock_setup_entry.mock_calls) == 1
assert mock_config.data[CONF_HOST] == "x.ui.direct"
assert mock_config.data[CONF_HOST] == DIRECT_CONNECT_DOMAIN
async def test_discovered_by_unifi_discovery_direct_connect_updated_but_not_using_direct_connect(
@ -518,3 +520,206 @@ async def test_discovered_by_unifi_discovery_partial(
"verify_ssl": False,
}
assert len(mock_setup_entry.mock_calls) == 1
async def test_discovered_by_unifi_discovery_direct_connect_on_different_interface(
hass: HomeAssistant, mock_nvr: NVR
) -> None:
"""Test a discovery from unifi-discovery from an alternate interface."""
mock_config = MockConfigEntry(
domain=DOMAIN,
data={
"host": DIRECT_CONNECT_DOMAIN,
"username": "test-username",
"password": "test-password",
"id": "UnifiProtect",
"port": 443,
"verify_ssl": True,
},
unique_id="FFFFFFAAAAAA",
)
mock_config.add_to_hass(hass)
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_ABORT
assert result["reason"] == "already_configured"
async def test_discovered_by_unifi_discovery_direct_connect_on_different_interface_ip_matches(
hass: HomeAssistant, mock_nvr: NVR
) -> None:
"""Test a discovery from unifi-discovery from an alternate interface when the ip matches."""
mock_config = MockConfigEntry(
domain=DOMAIN,
data={
"host": "127.0.0.1",
"username": "test-username",
"password": "test-password",
"id": "UnifiProtect",
"port": 443,
"verify_ssl": True,
},
unique_id="FFFFFFAAAAAA",
)
mock_config.add_to_hass(hass)
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_ABORT
assert result["reason"] == "already_configured"
async def test_discovered_by_unifi_discovery_direct_connect_on_different_interface_resolver(
hass: HomeAssistant, mock_nvr: NVR
) -> None:
"""Test a discovery from unifi-discovery from an alternate interface when direct connect domain resolves to host ip."""
mock_config = MockConfigEntry(
domain=DOMAIN,
data={
"host": "y.ui.direct",
"username": "test-username",
"password": "test-password",
"id": "UnifiProtect",
"port": 443,
"verify_ssl": True,
},
unique_id="FFFFFFAAAAAA",
)
mock_config.add_to_hass(hass)
other_ip_dict = UNIFI_DISCOVERY_DICT.copy()
other_ip_dict["source_ip"] = "127.0.0.1"
other_ip_dict["direct_connect_domain"] = "nomatchsameip.ui.direct"
with _patch_discovery(), patch.object(
hass.loop,
"getaddrinfo",
return_value=[(socket.AF_INET, None, None, None, ("127.0.0.1", 443))],
):
result = await hass.config_entries.flow.async_init(
DOMAIN,
context={"source": config_entries.SOURCE_DISCOVERY},
data=other_ip_dict,
)
await hass.async_block_till_done()
assert result["type"] == RESULT_TYPE_ABORT
assert result["reason"] == "already_configured"
async def test_discovered_by_unifi_discovery_direct_connect_on_different_interface_resolver_fails(
hass: HomeAssistant, mock_nvr: NVR
) -> None:
"""Test we can still configure if the resolver fails."""
mock_config = MockConfigEntry(
domain=DOMAIN,
data={
"host": "y.ui.direct",
"username": "test-username",
"password": "test-password",
"id": "UnifiProtect",
"port": 443,
"verify_ssl": True,
},
unique_id="FFFFFFAAAAAA",
)
mock_config.add_to_hass(hass)
other_ip_dict = UNIFI_DISCOVERY_DICT.copy()
other_ip_dict["source_ip"] = "127.0.0.2"
other_ip_dict["direct_connect_domain"] = "nomatchsameip.ui.direct"
with _patch_discovery(), patch.object(
hass.loop, "getaddrinfo", side_effect=OSError
):
result = await hass.config_entries.flow.async_init(
DOMAIN,
context={"source": config_entries.SOURCE_DISCOVERY},
data=other_ip_dict,
)
await hass.async_block_till_done()
assert result["type"] == RESULT_TYPE_FORM
assert result["step_id"] == "discovery_confirm"
flows = hass.config_entries.flow.async_progress_by_handler(DOMAIN)
assert flows[0]["context"]["title_placeholders"] == {
"ip_address": "127.0.0.2",
"name": "unvr",
}
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": "nomatchsameip.ui.direct",
"username": "test-username",
"password": "test-password",
"id": "UnifiProtect",
"port": 443,
"verify_ssl": True,
}
assert len(mock_setup_entry.mock_calls) == 1
async def test_discovered_by_unifi_discovery_direct_connect_on_different_interface_resolver_no_result(
hass: HomeAssistant, mock_nvr: NVR
) -> None:
"""Test a discovery from unifi-discovery from an alternate interface when direct connect domain resolve has no result."""
mock_config = MockConfigEntry(
domain=DOMAIN,
data={
"host": "y.ui.direct",
"username": "test-username",
"password": "test-password",
"id": "UnifiProtect",
"port": 443,
"verify_ssl": True,
},
unique_id="FFFFFFAAAAAA",
)
mock_config.add_to_hass(hass)
other_ip_dict = UNIFI_DISCOVERY_DICT.copy()
other_ip_dict["source_ip"] = "127.0.0.2"
other_ip_dict["direct_connect_domain"] = "y.ui.direct"
with _patch_discovery(), patch.object(hass.loop, "getaddrinfo", return_value=[]):
result = await hass.config_entries.flow.async_init(
DOMAIN,
context={"source": config_entries.SOURCE_DISCOVERY},
data=other_ip_dict,
)
await hass.async_block_till_done()
assert result["type"] == RESULT_TYPE_ABORT
assert result["reason"] == "already_configured"