Files
core/homeassistant/components/unifiprotect/entity.py
T

606 lines
23 KiB
Python

"""Shared Entity definition for UniFi Protect Integration."""
from collections.abc import Callable, Coroutine, Sequence
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from functools import partial
import logging
from operator import attrgetter
from typing import TYPE_CHECKING, Any, Generic, TypeVar, override
from uiprotect import make_enabled_getter, make_required_getter, make_value_getter
from uiprotect.data import (
NVR,
DeviceState,
Event,
ModelType,
ProtectAdoptableDeviceModel,
PublicDeviceModel,
SmartDetectObjectType,
StateType,
)
from uiprotect.data.public_devices import PublicSensor, SensorFeatureCapability
from homeassistant.const import Platform
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.helpers.device_registry import DeviceInfo
from homeassistant.helpers.entity import Entity, EntityDescription
from .const import (
ATTR_EVENT_ID,
ATTR_EVENT_SCORE,
ATTR_SMART_DETECT_TYPES,
DEFAULT_ATTRIBUTION,
DEFAULT_BRAND,
DOMAIN,
)
from .data import ProtectData, ProtectDeviceType
_LOGGER = logging.getLogger(__name__)
T = TypeVar("T", bound=ProtectAdoptableDeviceModel | NVR)
class PermRequired(int, Enum):
"""Type of permission level required for entity."""
NO_WRITE = 1
WRITE = 2
DELETE = 3
@callback
def _async_capability_supported(
data: ProtectData,
device: ProtectAdoptableDeviceModel,
description: ProtectEntityDescription,
) -> bool:
"""Whether the device advertises the description's required sensor capability."""
if (capability := description.ufp_capability) is None:
return True
public = data.async_get_public_device(device)
if not isinstance(public, PublicSensor) or not public.has_feature_flags:
return True
return public.supports(capability)
@callback
def async_remove_unsupported_sense_entities(
hass: HomeAssistant,
platform: Platform,
data: ProtectData,
descs: Sequence[ProtectEntityDescription],
) -> None:
"""Remove registry entries for sense entities the device cannot support.
Only acts when a public capability map is present (newer firmware); a console
upgrade then drops the never-functional entities created before the map existed.
"""
entity_registry = er.async_get(hass)
for device in data.get_by_types({ModelType.SENSOR}):
for description in descs:
if description.ufp_capability is None or _async_capability_supported(
data, device, description
):
continue
if entity_id := entity_registry.async_get_entity_id(
platform, DOMAIN, f"{device.mac}_{description.key}"
):
entity_registry.async_remove(entity_id)
@callback
def _async_device_entities(
data: ProtectData,
klass: type[BaseProtectEntity],
model_type: ModelType,
descs: Sequence[ProtectEntityDescription],
unadopted_descs: Sequence[ProtectEntityDescription] | None = None,
ufp_device: ProtectAdoptableDeviceModel | None = None,
) -> list[BaseProtectEntity]:
if not descs and not unadopted_descs:
return []
entities: list[BaseProtectEntity] = []
devices = (
[ufp_device]
if ufp_device is not None
else data.get_by_types({model_type}, ignore_unadopted=False)
)
auth_user = data.api.bootstrap.auth_user
for device in devices:
if TYPE_CHECKING:
assert isinstance(device, ProtectAdoptableDeviceModel)
if not device.is_adopted_by_us:
if unadopted_descs:
for description in unadopted_descs:
entities.append(
klass(
data,
device=device,
description=description,
)
)
_LOGGER.debug(
"Adding %s entity %s for %s",
klass.__name__,
description.key,
device.display_name,
)
continue
can_write = device.can_write(auth_user)
for description in descs:
if (perms := description.ufp_perm) is not None:
if perms is PermRequired.WRITE and not can_write:
continue
if perms is PermRequired.NO_WRITE and can_write:
continue
if perms is PermRequired.DELETE and not device.can_delete(auth_user):
continue
if not description.has_required(device):
continue
if not _async_capability_supported(data, device, description):
continue
entities.append(
klass(
data,
device=device,
description=description,
)
)
_LOGGER.debug(
"Adding %s entity %s for %s",
klass.__name__,
description.key,
device.display_name,
)
return entities
_ALL_MODEL_TYPES = (
ModelType.CAMERA,
ModelType.LIGHT,
ModelType.SENSOR,
ModelType.VIEWPORT,
ModelType.CHIME,
)
@callback
def _combine_model_descs(
model_type: ModelType,
model_descriptions: dict[ModelType, Sequence[ProtectEntityDescription]] | None,
all_descs: Sequence[ProtectEntityDescription] | None,
) -> list[ProtectEntityDescription]:
"""Combine all the descriptions with descriptions a model type."""
descs: list[ProtectEntityDescription] = list(all_descs) if all_descs else []
if model_descriptions and (model_descs := model_descriptions.get(model_type)):
descs.extend(model_descs)
return descs
@callback
def async_all_device_entities(
data: ProtectData,
klass: type[BaseProtectEntity],
model_descriptions: dict[ModelType, Sequence[ProtectEntityDescription]]
| None = None,
all_descs: Sequence[ProtectEntityDescription] | None = None,
unadopted_descs: list[ProtectEntityDescription] | None = None,
ufp_device: ProtectAdoptableDeviceModel | None = None,
) -> list[BaseProtectEntity]:
"""Generate a list of all the device entities."""
if ufp_device is None:
entities: list[BaseProtectEntity] = []
for model_type in _ALL_MODEL_TYPES:
descs = _combine_model_descs(model_type, model_descriptions, all_descs)
entities.extend(
_async_device_entities(data, klass, model_type, descs, unadopted_descs)
)
return entities
device_model_type = ufp_device.model
assert device_model_type is not None
# Runtime adoption must honor the same model-type allowlist as initial setup,
# so unsupported devices (e.g. AI Port) get no entities when adopted live.
if device_model_type not in _ALL_MODEL_TYPES:
return []
descs = _combine_model_descs(device_model_type, model_descriptions, all_descs)
return _async_device_entities(
data, klass, device_model_type, descs, unadopted_descs, ufp_device
)
class BaseProtectEntity(Entity):
"""Base class for UniFi protect entities."""
device: ProtectDeviceType
_attr_should_poll = False
_attr_attribution = DEFAULT_ATTRIBUTION
_state_attrs: tuple[str, ...] = ("_attr_available",)
_attr_has_entity_name = True
_async_get_ufp_enabled: Callable[[ProtectAdoptableDeviceModel], bool] | None = None
_async_get_ufp_public_enabled: Callable[[PublicDeviceModel], bool] | None = None
# Cached public-API object for descriptions migrated to the public path
# (set ``ufp_public_value``); ``None`` until primed/refreshed.
_ufp_public_obj: PublicDeviceModel | None = None
_ufp_uses_public: bool = False
# Values derived from the public events websocket (detection booleans,
# public event entities) additionally require that websocket to be healthy.
_ufp_requires_events_ws: bool = False
def __init__(
self,
data: ProtectData,
device: ProtectDeviceType,
description: EntityDescription | None = None,
) -> None:
"""Initialize the entity."""
super().__init__()
self.data = data
self.device = device
if description is None:
self._attr_unique_id = self.device.mac
self._attr_name = None
else:
self.entity_description = description
self._attr_unique_id = f"{self.device.mac}_{description.key}"
if isinstance(description, ProtectEntityDescription):
self._async_get_ufp_enabled = description.get_ufp_enabled
self._async_get_ufp_public_enabled = description.ufp_public_enabled_fn
self._async_set_device_info()
self._state_getters = tuple(
partial(attrgetter(attr), self) for attr in self._state_attrs
)
async def async_update(self) -> None:
"""Update the entity.
Only used by the generic entity update service.
"""
await self.data.async_refresh()
@callback
def _async_set_device_info(self) -> None:
"""Set device info."""
@callback
def _async_update_device_from_protect(self, device: ProtectDeviceType) -> None:
"""Update Entity object from Protect device."""
was_available = self._attr_available
if last_updated_success := self.data.last_update_success:
self.device = device
if self._ufp_uses_public:
# Migrated entities are fully public: availability tracks the public
# websocket health and the public object's state (CONNECTED only;
# CONNECTING/DISCONNECTED/UNKNOWN and a missing object read as
# unavailable), independent of the private connection. Values fed by
# the events websocket also require it to be healthy — the devices
# websocket keeps the device state fresh, but only the events stream
# carries the detections. An optional ``ufp_public_enabled_fn`` gate
# then mirrors ``ufp_enabled`` against the public object (e.g. a
# sensor feature toggled off).
public_obj = self._ufp_public_obj
if (
self.data.last_public_update_success
and (
not self._ufp_requires_events_ws
or self.data.last_events_update_success
)
and public_obj is not None
and public_obj.state is DeviceState.CONNECTED
):
get_public_enabled = self._async_get_ufp_public_enabled
available = get_public_enabled is None or get_public_enabled(public_obj)
else:
available = False
elif device.model is ModelType.NVR:
available = last_updated_success
else:
if TYPE_CHECKING:
assert isinstance(device, ProtectAdoptableDeviceModel)
connected = device.state is StateType.CONNECTED or (
not device.is_adopted_by_us and device.can_adopt
)
async_get_ufp_enabled = self._async_get_ufp_enabled
enabled = not async_get_ufp_enabled or async_get_ufp_enabled(device)
available = last_updated_success and connected and enabled
if available != was_available:
self._attr_available = available
@callback
def _async_updated_event(self, device: ProtectDeviceType) -> None:
"""When device is updated from Protect."""
previous_attrs = [getter() for getter in self._state_getters]
self._async_update_device_from_protect(device)
changed = False
for idx, getter in enumerate(self._state_getters):
if previous_attrs[idx] != getter():
changed = True
break
if changed:
if _LOGGER.isEnabledFor(logging.DEBUG):
_LOGGER.debug(
"Updating state [%s] %s -> %s",
self.entity_id,
previous_attrs,
tuple(getter() for getter in self._state_getters),
)
self.async_write_ha_state()
@callback
def _async_public_updated(self, obj: PublicDeviceModel | None) -> None:
"""Handle a public devices WS update for a migrated value.
``obj`` is the refreshed public object from a WS message; ``None`` when
there is no object to pass (a websocket state change, a delete event,
or a frame the library could not merge). The object is then re-read
from the bootstrap: a deleted device reads as missing (the entity goes
unavailable), and after a reconnect a value that changed during the
outage is picked up.
"""
self._ufp_public_obj = (
obj if obj is not None else self.data.async_get_public_device(self.device)
)
self._async_updated_event(self.device)
@override
async def async_added_to_hass(self) -> None:
"""When entity is added to hass."""
await super().async_added_to_hass()
self.async_on_remove(
self.data.async_subscribe(self.device.mac, self._async_updated_event)
)
# Not every entity carries an entity_description (e.g. cameras), so getattr.
description = getattr(self, "entity_description", None)
if isinstance(description, ProtectEntityDescription):
if (
description.ufp_public_value is not None
or description.ufp_public_value_fn is not None
):
self._ufp_uses_public = True
if description.ufp_event_driven:
self._ufp_requires_events_ws = True
# ``_ufp_uses_public`` may also be declared as a class attribute by
# entities driven by the public API without a migrated value (the
# public event entities).
if self._ufp_uses_public:
self._ufp_public_obj = self.data.async_get_public_device(self.device)
self.async_on_remove(
self.data.async_subscribe_public(
self.device.mac, self._async_public_updated
)
)
self._async_update_device_from_protect(self.device)
class ProtectIsOnEntity(BaseProtectEntity):
"""Base class for entities with is_on property."""
_state_attrs: tuple[str, ...] = ("_attr_available", "_attr_is_on")
_attr_is_on: bool | None
entity_description: ProtectEntityDescription
@override
def _async_update_device_from_protect(
self, device: ProtectAdoptableDeviceModel | NVR
) -> None:
super()._async_update_device_from_protect(device)
was_on = self._attr_is_on
value = self.entity_description.get_value(device, self._ufp_public_obj)
if was_on != (is_on := value is True):
self._attr_is_on = is_on
class ProtectDeviceEntity(BaseProtectEntity):
"""Base class for UniFi protect entities."""
@callback
@override
def _async_set_device_info(self) -> None:
self._attr_device_info = DeviceInfo(
name=self.device.display_name,
manufacturer=DEFAULT_BRAND,
model=self.device.market_name or self.device.type,
model_id=self.device.type,
via_device_id=self.data.nvr_device_id,
sw_version=self.device.firmware_version,
connections={(dr.CONNECTION_NETWORK_MAC, self.device.mac)},
configuration_url=self.device.protect_url,
)
class ProtectNVREntity(BaseProtectEntity):
"""Base class for unifi protect entities."""
device: NVR
@callback
@override
def _async_set_device_info(self) -> None:
self._attr_device_info = DeviceInfo(
connections={(dr.CONNECTION_NETWORK_MAC, self.device.mac)},
identifiers={(DOMAIN, self.device.mac)},
manufacturer=DEFAULT_BRAND,
name=self.device.display_name,
model=self.device.market_name or self.device.type,
model_id=self.device.type,
sw_version=str(self.device.version),
configuration_url=self.device.api.base_url,
)
class EventEntityMixin(ProtectDeviceEntity):
"""Adds motion event attributes to sensor."""
entity_description: ProtectEventMixin
_unrecorded_attributes = frozenset(
{ATTR_EVENT_ID, ATTR_EVENT_SCORE, ATTR_SMART_DETECT_TYPES}
)
_event: Event | None = None
_event_end: datetime | None = None
@callback
def _set_event_done(self) -> None:
"""Clear the event and state."""
@callback
def _set_event_attrs(self, event: Event) -> None:
"""Set event attrs."""
self._attr_extra_state_attributes = {
ATTR_EVENT_ID: event.id,
ATTR_EVENT_SCORE: event.score,
}
@callback
def _async_event_with_immediate_end(self) -> None:
# If the event is so short that the detection is received
# in the same message as the end of the event we need to write
# state and than clear the event and write state again.
self.async_write_ha_state()
self._set_event_done()
self.async_write_ha_state()
@callback
def _event_already_ended(
self, prev_event: Event | None, prev_event_end: datetime | None
) -> bool:
"""Determine if the event has already ended.
The event_end time is passed because the prev_event and event object
may be the same object, and the uiprotect code will mutate the
event object so we need to check the datetime object that was
saved from the last time the entity was updated.
"""
return bool(
(event := self._event)
and event.end
and prev_event
and prev_event_end
and prev_event.id == event.id
)
@dataclass(frozen=True, kw_only=True)
class ProtectEntityDescription(EntityDescription, Generic[T]): # noqa: UP046
"""Base class for protect entity descriptions."""
ufp_required_field: str | None = None
ufp_value: str | None = None
ufp_value_fn: Callable[[T], Any] | None = None
ufp_public_value: str | None = None
# Callable variant of ``ufp_public_value`` for public values needing a transform.
ufp_public_value_fn: Callable[[PublicDeviceModel], Any] | None = None
# True when the public value is derived from the events websocket (the
# detection booleans); availability then also tracks that websocket.
ufp_event_driven: bool = False
ufp_enabled: str | None = None
# Public counterpart of ``ufp_enabled``; a callable because public enablement
# is often compound (e.g. mount type plus a settings flag).
ufp_public_enabled_fn: Callable[[PublicDeviceModel], bool] | None = None
# Sensor capability required to create the entity, checked against the public
# capability map. Without a capability map (older firmware) every description
# is created, matching the pre-capability behavior.
ufp_capability: SensorFeatureCapability | None = None
ufp_perm: PermRequired | None = None
# The below are set in __post_init__
has_required: Callable[[T], bool] = bool
get_ufp_enabled: Callable[[T], bool] | None = None
get_ufp_public_value: Callable[[PublicDeviceModel], Any] | None = None
def get_ufp_value(self, obj: T) -> Any:
"""Return value from UniFi Protect device; overridden in __post_init__."""
# ufp_value or ufp_value_fn are required, the
# RuntimeError is to catch any issues in the code
# with new descriptions.
raise RuntimeError( # pragma: no cover
f"`ufp_value` or `ufp_value_fn` is required for {self}"
)
def get_value(self, obj: T, public_obj: PublicDeviceModel | None = None) -> Any:
"""Return the value, reading from the public object when migrated.
A migrated description sets ``ufp_public_value`` (or ``ufp_public_value_fn``)
and drops the private ``ufp_value``: the value comes only from the public
object, or ``None`` when it is absent (the entity is then marked
unavailable).
"""
if (fn := self.ufp_public_value_fn) is not None:
return None if public_obj is None else fn(public_obj)
if (getter := self.get_ufp_public_value) is not None:
return None if public_obj is None else getter(public_obj)
return self.get_ufp_value(obj)
def __post_init__(self) -> None:
"""Override get_ufp_value, has_required, and get_ufp_enabled if required."""
_setter = partial(object.__setattr__, self)
if (ufp_value := self.ufp_value) is not None:
_setter("get_ufp_value", make_value_getter(ufp_value))
elif (ufp_value_fn := self.ufp_value_fn) is not None:
_setter("get_ufp_value", ufp_value_fn)
if (ufp_public_value := self.ufp_public_value) is not None:
_setter("get_ufp_public_value", make_value_getter(ufp_public_value))
if (ufp_enabled := self.ufp_enabled) is not None:
_setter("get_ufp_enabled", make_enabled_getter(ufp_enabled))
if (ufp_required_field := self.ufp_required_field) is not None:
_setter("has_required", make_required_getter(ufp_required_field))
@dataclass(frozen=True, kw_only=True)
class ProtectEventMixin(ProtectEntityDescription[T]):
"""Mixin for events."""
ufp_event_obj: str | None = None
ufp_obj_type: SmartDetectObjectType | None = None
def get_event_obj(self, obj: T) -> Event | None:
"""Return value from UniFi Protect device."""
return None
def has_matching_smart(self, event: Event) -> bool:
"""Determine if the detection type is a match."""
return (
not (obj_type := self.ufp_obj_type) or obj_type in event.smart_detect_types
)
@override
def __post_init__(self) -> None:
"""Override get_event_obj if ufp_event_obj is set."""
if (_ufp_event_obj := self.ufp_event_obj) is not None:
object.__setattr__(self, "get_event_obj", attrgetter(_ufp_event_obj))
super().__post_init__()
@dataclass(frozen=True, kw_only=True)
class ProtectSettableKeysMixin(ProtectEntityDescription[T]):
"""Mixin for settable values."""
ufp_set_method: str | None = None
ufp_set_method_fn: Callable[[T, Any], Coroutine[Any, Any, None]] | None = None
async def ufp_set(self, obj: T, value: Any) -> None:
"""Set value for UniFi Protect device."""
_LOGGER.debug("Setting %s to %s for %s", self.key, value, obj.display_name)
if self.ufp_set_method is not None:
await getattr(obj, self.ufp_set_method)(value)
elif self.ufp_set_method_fn is not None:
await self.ufp_set_method_fn(obj, value)