UniFi - Allow tracking of clients connected to third party APs (#34067)

* Allow disable wired bug work around

* Move small improvements from closed PR #34065

* Fix failing test

* Add new test

* Some extra logging

* Harmonize log outputs

* Add config flow string

* Fix Balloobs comments
This commit is contained in:
Robert Svensson 2020-04-17 08:39:01 +02:00 committed by GitHub
parent 2326a2941e
commit 465eeab553
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
11 changed files with 145 additions and 64 deletions

View File

@ -37,6 +37,7 @@
"device_tracker": {
"data": {
"detection_time": "Time in seconds from last seen until considered away",
"ignore_wired_bug": "Disable UniFi wired bug logic",
"ssid_filter": "Select SSIDs to track wireless clients on",
"track_clients": "Track network clients",
"track_devices": "Track network devices (Ubiquiti devices)",

View File

@ -7,7 +7,7 @@ import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
from .config_flow import get_controller_id_from_config_entry
from .const import ATTR_MANUFACTURER, DOMAIN, UNIFI_WIRELESS_CLIENTS
from .const import ATTR_MANUFACTURER, DOMAIN, LOGGER, UNIFI_WIRELESS_CLIENTS
from .controller import UniFiController
SAVE_DELAY = 10
@ -42,6 +42,8 @@ async def async_setup_entry(hass, config_entry):
hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, controller.shutdown)
LOGGER.debug("UniFi config options %s", config_entry.options)
if controller.mac is None:
return True

View File

@ -19,6 +19,7 @@ from .const import (
CONF_BLOCK_CLIENT,
CONF_CONTROLLER,
CONF_DETECTION_TIME,
CONF_IGNORE_WIRED_BUG,
CONF_POE_CLIENTS,
CONF_SITE_ID,
CONF_SSID_FILTER,
@ -216,6 +217,10 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
self.controller.option_detection_time.total_seconds()
),
): int,
vol.Optional(
CONF_IGNORE_WIRED_BUG,
default=self.controller.option_ignore_wired_bug,
): bool,
}
),
)

View File

@ -14,6 +14,7 @@ UNIFI_WIRELESS_CLIENTS = "unifi_wireless_clients"
CONF_ALLOW_BANDWIDTH_SENSORS = "allow_bandwidth_sensors"
CONF_BLOCK_CLIENT = "block_client"
CONF_DETECTION_TIME = "detection_time"
CONF_IGNORE_WIRED_BUG = "ignore_wired_bug"
CONF_POE_CLIENTS = "poe_clients"
CONF_TRACK_CLIENTS = "track_clients"
CONF_TRACK_DEVICES = "track_devices"
@ -21,6 +22,7 @@ CONF_TRACK_WIRED_CLIENTS = "track_wired_clients"
CONF_SSID_FILTER = "ssid_filter"
DEFAULT_ALLOW_BANDWIDTH_SENSORS = False
DEFAULT_IGNORE_WIRED_BUG = False
DEFAULT_POE_CLIENTS = True
DEFAULT_TRACK_CLIENTS = True
DEFAULT_TRACK_DEVICES = True

View File

@ -31,6 +31,7 @@ from .const import (
CONF_BLOCK_CLIENT,
CONF_CONTROLLER,
CONF_DETECTION_TIME,
CONF_IGNORE_WIRED_BUG,
CONF_POE_CLIENTS,
CONF_SITE_ID,
CONF_SSID_FILTER,
@ -40,6 +41,7 @@ from .const import (
CONTROLLER_ID,
DEFAULT_ALLOW_BANDWIDTH_SENSORS,
DEFAULT_DETECTION_TIME,
DEFAULT_IGNORE_WIRED_BUG,
DEFAULT_POE_CLIENTS,
DEFAULT_TRACK_CLIENTS,
DEFAULT_TRACK_DEVICES,
@ -96,32 +98,20 @@ class UniFiController:
return self._site_role
@property
def option_allow_bandwidth_sensors(self):
"""Config entry option to allow bandwidth sensors."""
return self.config_entry.options.get(
CONF_ALLOW_BANDWIDTH_SENSORS, DEFAULT_ALLOW_BANDWIDTH_SENSORS
)
def mac(self):
"""Return the mac address of this controller."""
for client in self.api.clients.values():
if self.host == client.ip:
return client.mac
return None
@property
def option_block_clients(self):
"""Config entry option with list of clients to control network access."""
return self.config_entry.options.get(CONF_BLOCK_CLIENT, [])
@property
def option_poe_clients(self):
"""Config entry option to control poe clients."""
return self.config_entry.options.get(CONF_POE_CLIENTS, DEFAULT_POE_CLIENTS)
# Device tracker options
@property
def option_track_clients(self):
"""Config entry option to not track clients."""
return self.config_entry.options.get(CONF_TRACK_CLIENTS, DEFAULT_TRACK_CLIENTS)
@property
def option_track_devices(self):
"""Config entry option to not track devices."""
return self.config_entry.options.get(CONF_TRACK_DEVICES, DEFAULT_TRACK_DEVICES)
@property
def option_track_wired_clients(self):
"""Config entry option to not track wired clients."""
@ -129,6 +119,16 @@ class UniFiController:
CONF_TRACK_WIRED_CLIENTS, DEFAULT_TRACK_WIRED_CLIENTS
)
@property
def option_track_devices(self):
"""Config entry option to not track devices."""
return self.config_entry.options.get(CONF_TRACK_DEVICES, DEFAULT_TRACK_DEVICES)
@property
def option_ssid_filter(self):
"""Config entry option listing what SSIDs are being used to track clients."""
return self.config_entry.options.get(CONF_SSID_FILTER, [])
@property
def option_detection_time(self):
"""Config entry option defining number of seconds from last seen to away."""
@ -139,17 +139,32 @@ class UniFiController:
)
@property
def option_ssid_filter(self):
"""Config entry option listing what SSIDs are being used to track clients."""
return self.config_entry.options.get(CONF_SSID_FILTER, [])
def option_ignore_wired_bug(self):
"""Config entry option to ignore wired bug."""
return self.config_entry.options.get(
CONF_IGNORE_WIRED_BUG, DEFAULT_IGNORE_WIRED_BUG
)
# Client control options
@property
def mac(self):
"""Return the mac address of this controller."""
for client in self.api.clients.values():
if self.host == client.ip:
return client.mac
return None
def option_poe_clients(self):
"""Config entry option to control poe clients."""
return self.config_entry.options.get(CONF_POE_CLIENTS, DEFAULT_POE_CLIENTS)
@property
def option_block_clients(self):
"""Config entry option with list of clients to control network access."""
return self.config_entry.options.get(CONF_BLOCK_CLIENT, [])
# Statistics sensor options
@property
def option_allow_bandwidth_sensors(self):
"""Config entry option to allow bandwidth sensors."""
return self.config_entry.options.get(
CONF_ALLOW_BANDWIDTH_SENSORS, DEFAULT_ALLOW_BANDWIDTH_SENSORS
)
@callback
def async_unifi_signalling_callback(self, signal, data):

View File

@ -67,6 +67,9 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
client = controller.api.clients_all[mac]
controller.api.clients.process_raw([client.raw])
LOGGER.debug(
"Restore disconnected client %s (%s)", entity.entity_id, client.mac,
)
@callback
def items_added():
@ -130,20 +133,21 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
remove.add(mac)
if option_ssid_filter != controller.option_ssid_filter:
option_ssid_filter = controller.option_ssid_filter
update = True
for mac, entity in tracked.items():
if (
isinstance(entity, UniFiClientTracker)
and not entity.is_wired
and entity.client.essid not in option_ssid_filter
):
remove.add(mac)
if controller.option_ssid_filter:
for mac, entity in tracked.items():
if (
isinstance(entity, UniFiClientTracker)
and not entity.is_wired
and entity.client.essid not in controller.option_ssid_filter
):
remove.add(mac)
option_track_clients = controller.option_track_clients
option_track_devices = controller.option_track_devices
option_track_wired_clients = controller.option_track_wired_clients
option_ssid_filter = controller.option_ssid_filter
remove_entities(controller, remove, tracked, entity_registry)
@ -319,13 +323,12 @@ class UniFiDeviceTracker(ScannerEntity):
"""Set up tracked device."""
self.device = device
self.controller = controller
self.listeners = []
async def async_added_to_hass(self):
"""Subscribe to device events."""
LOGGER.debug("New UniFi device tracker %s (%s)", self.name, self.device.mac)
LOGGER.debug("New device %s (%s)", self.entity_id, self.device.mac)
self.device.register_callback(self.async_update_callback)
self.listeners.append(
self.async_on_remove(
async_dispatcher_connect(
self.hass, self.controller.signal_reachable, self.async_update_callback
)
@ -334,13 +337,11 @@ class UniFiDeviceTracker(ScannerEntity):
async def async_will_remove_from_hass(self) -> None:
"""Disconnect device object when removed."""
self.device.remove_callback(self.async_update_callback)
for unsub_dispatcher in self.listeners:
unsub_dispatcher()
@callback
def async_update_callback(self):
"""Update the sensor's state."""
LOGGER.debug("Updating UniFi tracked device %s", self.entity_id)
LOGGER.debug("Updating device %s (%s)", self.entity_id, self.device.mac)
self.async_write_ha_state()

View File

@ -118,7 +118,7 @@ class UniFiRxBandwidthSensor(UniFiClient):
@property
def state(self):
"""Return the state of the sensor."""
if self.is_wired:
if self._is_wired:
return self.client.wired_rx_bytes / 1000000
return self.client.raw.get("rx_bytes", 0) / 1000000
@ -145,7 +145,7 @@ class UniFiTxBandwidthSensor(UniFiRxBandwidthSensor):
@property
def state(self):
"""Return the state of the sensor."""
if self.is_wired:
if self._is_wired:
return self.client.wired_tx_bytes / 1000000
return self.client.raw.get("tx_bytes", 0) / 1000000

View File

@ -29,6 +29,7 @@
"device_tracker": {
"data": {
"detection_time": "Time in seconds from last seen until considered away",
"ignore_wired_bug": "Disable UniFi wired bug logic",
"ssid_filter": "Select SSIDs to track wireless clients on",
"track_clients": "Track network clients",
"track_devices": "Track network devices (Ubiquiti devices)",
@ -55,4 +56,4 @@
}
}
}
}
}

View File

@ -39,18 +39,17 @@ class UniFiClient(Entity):
"""Set up client."""
self.client = client
self.controller = controller
self.listeners = []
self.is_wired = self.client.mac not in controller.wireless_clients
self._is_wired = self.client.mac not in controller.wireless_clients
self.is_blocked = self.client.blocked
self.wired_connection = None
self.wireless_connection = None
async def async_added_to_hass(self) -> None:
"""Client entity created."""
LOGGER.debug("New UniFi client %s (%s)", self.name, self.client.mac)
LOGGER.debug("New client %s (%s)", self.entity_id, self.client.mac)
self.client.register_callback(self.async_update_callback)
self.listeners.append(
self.async_on_remove(
async_dispatcher_connect(
self.hass, self.controller.signal_reachable, self.async_update_callback
)
@ -59,17 +58,14 @@ class UniFiClient(Entity):
async def async_will_remove_from_hass(self) -> None:
"""Disconnect client object when removed."""
self.client.remove_callback(self.async_update_callback)
for unsub_dispatcher in self.listeners:
unsub_dispatcher()
@callback
def async_update_callback(self) -> None:
"""Update the clients state."""
if self.is_wired and self.client.mac in self.controller.wireless_clients:
self.is_wired = False
if self._is_wired and self.client.mac in self.controller.wireless_clients:
self._is_wired = False
if self.client.last_updated == SOURCE_EVENT:
if self.client.event.event in WIRELESS_CLIENT:
self.wireless_connection = self.client.event.event in (
WIRELESS_CLIENT_CONNECTED,
@ -84,9 +80,19 @@ class UniFiClient(Entity):
elif self.client.event.event in CLIENT_BLOCKED + CLIENT_UNBLOCKED:
self.is_blocked = self.client.event.event in CLIENT_BLOCKED
LOGGER.debug("Updating client %s %s", self.entity_id, self.client.mac)
LOGGER.debug("Updating client %s (%s)", self.entity_id, self.client.mac)
self.async_write_ha_state()
@property
def is_wired(self):
"""Return if the client is wired.
Allows disabling logic to keep track of clients affected by UniFi wired bug marking wireless devices as wired. This is useful when running a network not only containing UniFi APs.
"""
if self.controller.option_ignore_wired_bug:
return self.client.is_wired
return self._is_wired
@property
def name(self) -> str:
"""Return the name of the client."""

View File

@ -11,6 +11,7 @@ from homeassistant.components.unifi.const import (
CONF_BLOCK_CLIENT,
CONF_CONTROLLER,
CONF_DETECTION_TIME,
CONF_IGNORE_WIRED_BUG,
CONF_POE_CLIENTS,
CONF_SITE_ID,
CONF_SSID_FILTER,
@ -338,9 +339,10 @@ async def test_option_flow(hass):
CONF_TRACK_CLIENTS: False,
CONF_TRACK_WIRED_CLIENTS: False,
CONF_TRACK_DEVICES: False,
CONF_DETECTION_TIME: 100,
CONF_SSID_FILTER: ["SSID 1"],
CONF_BLOCK_CLIENT: ["00:00:00:00:00:01"],
CONF_DETECTION_TIME: 100,
CONF_IGNORE_WIRED_BUG: False,
CONF_POE_CLIENTS: False,
CONF_BLOCK_CLIENT: ["00:00:00:00:00:01"],
CONF_ALLOW_BANDWIDTH_SENSORS: True,
}

View File

@ -11,6 +11,7 @@ from homeassistant.components import unifi
import homeassistant.components.device_tracker as device_tracker
from homeassistant.components.unifi.const import (
CONF_BLOCK_CLIENT,
CONF_IGNORE_WIRED_BUG,
CONF_SSID_FILTER,
CONF_TRACK_CLIENTS,
CONF_TRACK_DEVICES,
@ -376,12 +377,18 @@ async def test_option_track_devices(hass):
async def test_option_ssid_filter(hass):
"""Test the SSID filter works."""
controller = await setup_unifi_integration(
hass, options={CONF_SSID_FILTER: ["ssid"]}, clients_response=[CLIENT_3],
)
assert len(hass.states.async_entity_ids("device_tracker")) == 0
controller = await setup_unifi_integration(hass, clients_response=[CLIENT_3])
assert len(hass.states.async_entity_ids("device_tracker")) == 1
client_3 = hass.states.get("device_tracker.client_3")
assert client_3
# Set SSID filter
hass.config_entries.async_update_entry(
controller.config_entry, options={CONF_SSID_FILTER: ["ssid"]},
)
await hass.async_block_till_done()
# SSID filter active
client_3 = hass.states.get("device_tracker.client_3")
assert not client_3
@ -403,7 +410,6 @@ async def test_option_ssid_filter(hass):
controller.api.message_handler(event)
await hass.async_block_till_done()
# SSID no longer filtered
client_3 = hass.states.get("device_tracker.client_3")
assert client_3.state == "home"
@ -422,6 +428,7 @@ async def test_wireless_client_go_wired_issue(hass):
client_1 = hass.states.get("device_tracker.client_1")
assert client_1 is not None
assert client_1.state == "home"
assert client_1.attributes["is_wired"] is False
client_1_client["is_wired"] = True
client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
@ -431,6 +438,7 @@ async def test_wireless_client_go_wired_issue(hass):
client_1 = hass.states.get("device_tracker.client_1")
assert client_1.state == "home"
assert client_1.attributes["is_wired"] is False
with patch.object(
unifi.device_tracker.dt_util,
@ -443,6 +451,7 @@ async def test_wireless_client_go_wired_issue(hass):
client_1 = hass.states.get("device_tracker.client_1")
assert client_1.state == "not_home"
assert client_1.attributes["is_wired"] is False
client_1_client["is_wired"] = False
client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
@ -452,6 +461,43 @@ async def test_wireless_client_go_wired_issue(hass):
client_1 = hass.states.get("device_tracker.client_1")
assert client_1.state == "home"
assert client_1.attributes["is_wired"] is False
async def test_option_ignore_wired_bug(hass):
"""Test option to ignore wired bug."""
client_1_client = copy(CLIENT_1)
client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
controller = await setup_unifi_integration(
hass, options={CONF_IGNORE_WIRED_BUG: True}, clients_response=[client_1_client]
)
assert len(hass.states.async_entity_ids("device_tracker")) == 1
client_1 = hass.states.get("device_tracker.client_1")
assert client_1 is not None
assert client_1.state == "home"
assert client_1.attributes["is_wired"] is False
client_1_client["is_wired"] = True
client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
event = {"meta": {"message": "sta:sync"}, "data": [client_1_client]}
controller.api.message_handler(event)
await hass.async_block_till_done()
client_1 = hass.states.get("device_tracker.client_1")
assert client_1.state == "home"
assert client_1.attributes["is_wired"] is True
client_1_client["is_wired"] = False
client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
event = {"meta": {"message": "sta:sync"}, "data": [client_1_client]}
controller.api.message_handler(event)
await hass.async_block_till_done()
client_1 = hass.states.get("device_tracker.client_1")
assert client_1.state == "home"
assert client_1.attributes["is_wired"] is False
async def test_restoring_client(hass):