device_tracker.py

Go to the documentation of this file.

"""Track both clients and devices using UniFi Network."""
 
from __future__ import annotations
 
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from datetime import timedelta
import logging
from typing import Any
 
import aiounifi
from aiounifi.interfaces.api_handlers import ItemEvent
from aiounifi.interfaces.clients import Clients
from aiounifi.interfaces.devices import Devices
from aiounifi.models.api import ApiItemT
from aiounifi.models.client import Client
from aiounifi.models.device import Device
from aiounifi.models.event import Event, EventKey
from propcache import cached_property
 
from homeassistant.components.device_tracker import (
    DOMAIN as DEVICE_TRACKER_DOMAIN,
    ScannerEntity,
    ScannerEntityDescription,
)
from homeassistant.core import Event as core_Event, HomeAssistant, callback
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity_platform import AddEntitiesCallback
import homeassistant.helpers.entity_registry as er
import homeassistant.util.dt as dt_util
 
from . import UnifiConfigEntry
from .const import DOMAIN as UNIFI_DOMAIN
from .entity import (
    HandlerT,
    UnifiEntity,
    UnifiEntityDescription,
    async_device_available_fn,
)
from .hub import UnifiHub
 
LOGGER = logging.getLogger(__name__)
 
CLIENT_TRACKER = "client"
DEVICE_TRACKER = "device"
 
CLIENT_CONNECTED_ATTRIBUTES = [
    "_is_guest_by_uap",
    "ap_mac",
    "authorized",
    "essid",
    "ip",
    "is_11r",
    "is_guest",
    "note",
    "qos_policy_applied",
    "radio",
    "radio_proto",
    "vlan",
]
 
CLIENT_STATIC_ATTRIBUTES = [
    "mac",
    "name",
    "oui",
]
 
 
CLIENT_CONNECTED_ALL_ATTRIBUTES = CLIENT_CONNECTED_ATTRIBUTES + CLIENT_STATIC_ATTRIBUTES
 
WIRED_CONNECTION = (EventKey.WIRED_CLIENT_CONNECTED,)
WIRED_DISCONNECTION = (EventKey.WIRED_CLIENT_DISCONNECTED,)
WIRELESS_CONNECTION = (
    EventKey.WIRELESS_CLIENT_CONNECTED,
    EventKey.WIRELESS_CLIENT_ROAM,
    EventKey.WIRELESS_CLIENT_ROAM_RADIO,
    EventKey.WIRELESS_GUEST_CONNECTED,
    EventKey.WIRELESS_GUEST_ROAM,
    EventKey.WIRELESS_GUEST_ROAM_RADIO,
)
WIRELESS_DISCONNECTION = (
    EventKey.WIRELESS_CLIENT_DISCONNECTED,
    EventKey.WIRELESS_GUEST_DISCONNECTED,
)
 
 
@callback
def async_client_allowed_fn(hub: UnifiHub, obj_id: str) -> bool:
    """Check if client is allowed."""
    if obj_id in hub.config.option_supported_clients:
        return True
 
    if not hub.config.option_track_clients:
        return False
 
    client = hub.api.clients[obj_id]
    if client.mac not in hub.entity_loader.wireless_clients:
        if not hub.config.option_track_wired_clients:
            return False
 
    elif (
        client.essid
        and hub.config.option_ssid_filter
        and client.essid not in hub.config.option_ssid_filter
    ):
        return False
 
    return True
 
 
@callback
def async_client_is_connected_fn(hub: UnifiHub, obj_id: str) -> bool:
    """Check if device object is disabled."""
    client = hub.api.clients[obj_id]
 
    if hub.entity_loader.wireless_clients.is_wireless(client) and client.is_wired:
        if not hub.config.option_ignore_wired_bug:
            return False  # Wired bug in action
 
    if (
        not client.is_wired
        and client.essid
        and hub.config.option_ssid_filter
        and client.essid not in hub.config.option_ssid_filter
    ):
        return False
 
    if (
        dt_util.utcnow() - dt_util.utc_from_timestamp(client.last_seen or 0)
        > hub.config.option_detection_time
    ):
        return False
 
    return True
 
 
@callback
def async_device_heartbeat_timedelta_fn(hub: UnifiHub, obj_id: str) -> timedelta:
    """Check if device object is disabled."""
    device = hub.api.devices[obj_id]
    return timedelta(seconds=device.next_interval + 60)
 
 
@dataclass(frozen=True, kw_only=True)
class UnifiTrackerEntityDescription(
    UnifiEntityDescription[HandlerT, ApiItemT], ScannerEntityDescription
):
    """Class describing UniFi device tracker entity."""
 
    heartbeat_timedelta_fn: Callable[[UnifiHub, str], timedelta]
    ip_address_fn: Callable[[aiounifi.Controller, str], str | None]
    is_connected_fn: Callable[[UnifiHub, str], bool]
    hostname_fn: Callable[[aiounifi.Controller, str], str | None]
 
 
ENTITY_DESCRIPTIONS: tuple[UnifiTrackerEntityDescription, ...] = (
    UnifiTrackerEntityDescription[Clients, Client](
        key="Client device scanner",
        allowed_fn=async_client_allowed_fn,
        api_handler_fn=lambda api: api.clients,
        device_info_fn=lambda api, obj_id: None,
        event_is_on=set(WIRED_CONNECTION + WIRELESS_CONNECTION),
        event_to_subscribe=(
            WIRED_CONNECTION
            + WIRED_DISCONNECTION
            + WIRELESS_CONNECTION
            + WIRELESS_DISCONNECTION
        ),
        heartbeat_timedelta_fn=lambda hub, _: hub.config.option_detection_time,
        is_connected_fn=async_client_is_connected_fn,
        name_fn=lambda client: client.name or client.hostname,
        object_fn=lambda api, obj_id: api.clients[obj_id],
        unique_id_fn=lambda hub, obj_id: f"{hub.site}-{obj_id}",
        ip_address_fn=lambda api, obj_id: api.clients[obj_id].ip,
        hostname_fn=lambda api, obj_id: api.clients[obj_id].hostname,
    ),
    UnifiTrackerEntityDescription[Devices, Device](
        key="Device scanner",
        allowed_fn=lambda hub, obj_id: hub.config.option_track_devices,
        api_handler_fn=lambda api: api.devices,
        available_fn=async_device_available_fn,
        device_info_fn=lambda api, obj_id: None,
        heartbeat_timedelta_fn=async_device_heartbeat_timedelta_fn,
        is_connected_fn=lambda hub, obj_id: hub.api.devices[obj_id].state == 1,
        name_fn=lambda device: device.name or device.model,
        object_fn=lambda api, obj_id: api.devices[obj_id],
        unique_id_fn=lambda hub, obj_id: obj_id,
        ip_address_fn=lambda api, obj_id: api.devices[obj_id].ip,
        hostname_fn=lambda api, obj_id: None,
    ),
)
 
 
@callback
def async_update_unique_id(hass: HomeAssistant, config_entry: UnifiConfigEntry) -> None:
    """Normalize client unique ID to have a prefix rather than suffix.
 
    Introduced with release 2023.12.
    """
    hub = config_entry.runtime_data
    ent_reg = er.async_get(hass)
 
    @callback
    def update_unique_id(obj_id: str) -> None:
        """Rework unique ID."""
        new_unique_id = f"{hub.site}-{obj_id}"
        if ent_reg.async_get_entity_id(
            DEVICE_TRACKER_DOMAIN, UNIFI_DOMAIN, new_unique_id
        ):
            return
 
        unique_id = f"{obj_id}-{hub.site}"
        if entity_id := ent_reg.async_get_entity_id(
            DEVICE_TRACKER_DOMAIN, UNIFI_DOMAIN, unique_id
        ):
            ent_reg.async_update_entity(entity_id, new_unique_id=new_unique_id)
 
    for obj_id in list(hub.api.clients) + list(hub.api.clients_all):
        update_unique_id(obj_id)
 
 
async def async_setup_entry(
    hass: HomeAssistant,
    config_entry: UnifiConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Set up device tracker for UniFi Network integration."""
    async_update_unique_id(hass, config_entry)
    config_entry.runtime_data.entity_loader.register_platform(
        async_add_entities, UnifiScannerEntity, ENTITY_DESCRIPTIONS
    )
 
 
class UnifiScannerEntity(UnifiEntity[HandlerT, ApiItemT], ScannerEntity):
    """Representation of a UniFi scanner."""
 
    entity_description: UnifiTrackerEntityDescription
 
    _event_is_on: set[EventKey]
    _ignore_events: bool
    _is_connected: bool
 
    @callback
    def async_initiate_state(self) -> None:
        """Initiate entity state.
 
        Initiate is_connected.
        """
        description = self.entity_descriptionentity_description
        self._event_is_on_event_is_on = description.event_is_on or set()
        self._ignore_events_ignore_events = False
        self._is_connected_is_connected = description.is_connected_fn(self.hubhub, self._obj_id_obj_id)
        if self.is_connectedis_connected:
            self.hubhub.update_heartbeat(
                self.unique_idunique_idunique_id,
                dt_util.utcnow()
                + description.heartbeat_timedelta_fn(self.hubhub, self._obj_id_obj_id),
            )
 
    @property
    def is_connected(self) -> bool:
        """Return true if the device is connected to the network."""
        return self._is_connected_is_connected
 
    @property
    def hostname(self) -> str | None:
        """Return hostname of the device."""
        return self.entity_descriptionentity_description.hostname_fn(self.apiapi, self._obj_id_obj_id)
 
    @property
    def ip_address(self) -> str | None:
        """Return the primary ip address of the device."""
        return self.entity_descriptionentity_description.ip_address_fn(self.apiapi, self._obj_id_obj_id)
 
    @cached_property
    def mac_address(self) -> str:
        """Return the mac address of the device."""
        return self._obj_id_obj_id
 
    @cached_property
    def unique_id(self) -> str:
        """Return a unique ID."""
        return self._attr_unique_id_attr_unique_id
 
    @callback
    def _make_disconnected(self, *_: core_Event) -> None:
        """No heart beat by device."""
        self._is_connected_is_connected = False
        self.async_write_ha_stateasync_write_ha_state()
 
    @callback
    def async_update_state(self, event: ItemEvent, obj_id: str) -> None:
        """Update entity state.
 
        Remove heartbeat check if hub connection state has changed
         and entity is unavailable.
        Update is_connected.
        Schedule new heartbeat check if connected.
        """
        description = self.entity_descriptionentity_description
        hub = self.hubhub
 
        if event is ItemEvent.CHANGED:
            # Prioritize normal data updates over events
            self._ignore_events_ignore_events = True
 
        elif event is ItemEvent.ADDED and not self.availableavailable:
            # From unifi.entity.async_signal_reachable_callback
            # Controller connection state has changed and entity is unavailable
            # Cancel heartbeat
            hub.remove_heartbeat(self.unique_idunique_idunique_id)
            return
 
        obj_id = self._obj_id_obj_id
        if is_connected := description.is_connected_fn(hub, obj_id):
            self._is_connected_is_connected = is_connected
            self.hubhub.update_heartbeat(
                self.unique_idunique_idunique_id,
                dt_util.utcnow() + description.heartbeat_timedelta_fn(hub, obj_id),
            )
 
    @callback
    def async_event_callback(self, event: Event) -> None:
        """Event subscription callback."""
        obj_id = self._obj_id_obj_id
        if event.mac != obj_id or self._ignore_events_ignore_events:
            return
 
        hub = self.hubhub
        if event.key in self._event_is_on_event_is_on:
            hub.remove_heartbeat(self.unique_idunique_idunique_id)
            self._is_connected_is_connected = True
            self.async_write_ha_stateasync_write_ha_state()
            return
 
        hub.update_heartbeat(
            self.unique_idunique_idunique_id,
            dt_util.utcnow()
            + self.entity_descriptionentity_description.heartbeat_timedelta_fn(hub, obj_id),
        )
 
    async def async_added_to_hass(self) -> None:
        """Register callbacks."""
        await super().async_added_to_hass()
        self.async_on_removeasync_on_remove(
            async_dispatcher_connect(
                self.hasshass,
                f"{self.hub.signal_heartbeat_missed}_{self.unique_id}",
                self._make_disconnected_make_disconnected,
            )
        )
 
    async def async_will_remove_from_hass(self) -> None:
        """Disconnect object when removed."""
        await super().async_will_remove_from_hass()
        self.hubhub.remove_heartbeat(self.unique_idunique_idunique_id)
 
    @property
    def extra_state_attributes(self) -> Mapping[str, Any] | None:
        """Return the client state attributes."""
        if self.entity_descriptionentity_description.key != "Client device scanner":
            return None
 
        client = self.entity_descriptionentity_description.object_fn(self.apiapi, self._obj_id_obj_id)
        raw = client.raw
 
        attributes_to_check = CLIENT_STATIC_ATTRIBUTES
        if self.is_connectedis_connected:
            attributes_to_check = CLIENT_CONNECTED_ALL_ATTRIBUTES
 
        return {k: raw[k] for k in attributes_to_check if k in raw}