core/homeassistant/components/unifi/device_tracker.py

"""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_description
        self._event_is_on = description.event_is_on or set()
        self._ignore_events = False
        self._is_connected = description.is_connected_fn(self.hub, self._obj_id)
        if self.is_connected:
            self.hub.update_heartbeat(
                self.unique_id,
                dt_util.utcnow()
                + description.heartbeat_timedelta_fn(self.hub, self._obj_id),
            )

    @property
    def is_connected(self) -> bool:
        """Return true if the device is connected to the network."""
        return self._is_connected

    @property
    def hostname(self) -> str | None:
        """Return hostname of the device."""
        return self.entity_description.hostname_fn(self.api, self._obj_id)

    @property
    def ip_address(self) -> str | None:
        """Return the primary ip address of the device."""
        return self.entity_description.ip_address_fn(self.api, self._obj_id)

    @cached_property
    def mac_address(self) -> str:
        """Return the mac address of the device."""
        return self._obj_id

    @cached_property
    def unique_id(self) -> str:
        """Return a unique ID."""
        return self._attr_unique_id

    @callback
    def _make_disconnected(self, *_: core_Event) -> None:
        """No heart beat by device."""
        self._is_connected = False
        self.async_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_description
        hub = self.hub

        if event is ItemEvent.CHANGED:
            # Prioritize normal data updates over events
            self._ignore_events = True

        elif event is ItemEvent.ADDED and not self.available:
            # From unifi.entity.async_signal_reachable_callback
            # Controller connection state has changed and entity is unavailable
            # Cancel heartbeat
            hub.remove_heartbeat(self.unique_id)
            return

        obj_id = self._obj_id
        if is_connected := description.is_connected_fn(hub, obj_id):
            self._is_connected = is_connected
            self.hub.update_heartbeat(
                self.unique_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
        if event.mac != obj_id or self._ignore_events:
            return

        hub = self.hub
        if event.key in self._event_is_on:
            hub.remove_heartbeat(self.unique_id)
            self._is_connected = True
            self.async_write_ha_state()
            return

        hub.update_heartbeat(
            self.unique_id,
            dt_util.utcnow()
            + self.entity_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_remove(
            async_dispatcher_connect(
                self.hass,
                f"{self.hub.signal_heartbeat_missed}_{self.unique_id}",
                self._make_disconnected,
            )
        )

    async def async_will_remove_from_hass(self) -> None:
        """Disconnect object when removed."""
        await super().async_will_remove_from_hass()
        self.hub.remove_heartbeat(self.unique_id)

    @property
    def extra_state_attributes(self) -> Mapping[str, Any] | None:
        """Return the client state attributes."""
        if self.entity_description.key != "Client device scanner":
            return None

        client = self.entity_description.object_fn(self.api, self._obj_id)
        raw = client.raw

        attributes_to_check = CLIENT_STATIC_ATTRIBUTES
        if self.is_connected:
            attributes_to_check = CLIENT_CONNECTED_ALL_ATTRIBUTES

        return {k: raw[k] for k in attributes_to_check if k in raw}