"""Shared Entity definition for UniFi Protect Integration.""" from __future__ import annotations from collections import deque from collections.abc import Sequence from datetime import datetime, timedelta import hashlib import logging from random import SystemRandom from typing import Any, Final from urllib.parse import urlencode from pyunifiprotect.data import ( Camera, Event, Light, ModelType, ProtectAdoptableDeviceModel, Sensor, StateType, Viewer, ) from pyunifiprotect.data.nvr import NVR from homeassistant.core import callback import homeassistant.helpers.device_registry as dr from homeassistant.helpers.entity import DeviceInfo, Entity, EntityDescription from .const import ( ATTR_EVENT_SCORE, ATTR_EVENT_THUMB, DEFAULT_ATTRIBUTION, DEFAULT_BRAND, DOMAIN, ) from .data import ProtectData from .models import ProtectRequiredKeysMixin from .utils import get_nested_attr from .views import ThumbnailProxyView EVENT_UPDATE_TOKENS = "unifiprotect_update_tokens" TOKEN_CHANGE_INTERVAL: Final = timedelta(minutes=1) _RND: Final = SystemRandom() _LOGGER = logging.getLogger(__name__) @callback def _async_device_entities( data: ProtectData, klass: type[ProtectDeviceEntity], model_type: ModelType, descs: Sequence[ProtectRequiredKeysMixin], ) -> list[ProtectDeviceEntity]: if len(descs) == 0: return [] entities: list[ProtectDeviceEntity] = [] for device in data.get_by_types({model_type}): assert isinstance(device, (Camera, Light, Sensor, Viewer)) for description in descs: assert isinstance(description, EntityDescription) if description.ufp_required_field: required_field = get_nested_attr(device, description.ufp_required_field) if not required_field: continue entities.append( klass( data, device=device, description=description, ) ) _LOGGER.debug( "Adding %s entity %s for %s", klass.__name__, description.name, device.name, ) return entities @callback def async_all_device_entities( data: ProtectData, klass: type[ProtectDeviceEntity], camera_descs: Sequence[ProtectRequiredKeysMixin] | None = None, light_descs: Sequence[ProtectRequiredKeysMixin] | None = None, sense_descs: Sequence[ProtectRequiredKeysMixin] | None = None, viewer_descs: Sequence[ProtectRequiredKeysMixin] | None = None, all_descs: Sequence[ProtectRequiredKeysMixin] | None = None, ) -> list[ProtectDeviceEntity]: """Generate a list of all the device entities.""" all_descs = list(all_descs or []) camera_descs = list(camera_descs or []) + all_descs light_descs = list(light_descs or []) + all_descs sense_descs = list(sense_descs or []) + all_descs viewer_descs = list(viewer_descs or []) + all_descs return ( _async_device_entities(data, klass, ModelType.CAMERA, camera_descs) + _async_device_entities(data, klass, ModelType.LIGHT, light_descs) + _async_device_entities(data, klass, ModelType.SENSOR, sense_descs) + _async_device_entities(data, klass, ModelType.VIEWPORT, viewer_descs) ) class ProtectDeviceEntity(Entity): """Base class for UniFi protect entities.""" device: ProtectAdoptableDeviceModel _attr_should_poll = False def __init__( self, data: ProtectData, device: ProtectAdoptableDeviceModel, description: EntityDescription | None = None, ) -> None: """Initialize the entity.""" super().__init__() self.data: ProtectData = data self.device = device if description is None: self._attr_unique_id = f"{self.device.id}" self._attr_name = f"{self.device.name}" else: self.entity_description = description self._attr_unique_id = f"{self.device.id}_{description.key}" name = description.name or "" self._attr_name = f"{self.device.name} {name.title()}" self._attr_attribution = DEFAULT_ATTRIBUTION self._async_set_device_info() self._async_update_device_from_protect() 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: self._attr_device_info = DeviceInfo( name=self.device.name, manufacturer=DEFAULT_BRAND, model=self.device.type, via_device=(DOMAIN, self.data.api.bootstrap.nvr.mac), sw_version=self.device.firmware_version, connections={(dr.CONNECTION_NETWORK_MAC, self.device.mac)}, configuration_url=self.device.protect_url, ) @callback def _async_update_device_from_protect(self) -> None: """Update Entity object from Protect device.""" if self.data.last_update_success: assert self.device.model devices = getattr(self.data.api.bootstrap, f"{self.device.model.value}s") self.device = devices[self.device.id] self._attr_available = ( self.data.last_update_success and self.device.state == StateType.CONNECTED ) @callback def _async_updated_event(self) -> None: """Call back for incoming data.""" self._async_update_device_from_protect() self.async_write_ha_state() 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_device_id( self.device.id, self._async_updated_event ) ) class ProtectNVREntity(ProtectDeviceEntity): """Base class for unifi protect entities.""" # separate subclass on purpose device: NVR # type: ignore[assignment] def __init__( self, entry: ProtectData, device: NVR, description: EntityDescription | None = None, ) -> None: """Initialize the entity.""" super().__init__(entry, device, description) # type: ignore[arg-type] @callback 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.name, model=self.device.type, sw_version=str(self.device.version), configuration_url=self.device.api.base_url, ) @callback def _async_update_device_from_protect(self) -> None: if self.data.last_update_success: self.device = self.data.api.bootstrap.nvr self._attr_available = self.data.last_update_success class AccessTokenMixin(ProtectDeviceEntity): """Adds access_token attribute and provides access tokens for use for anonymous views.""" @property def access_tokens(self) -> deque[str]: """Get valid access_tokens for current entity.""" return self.data.async_get_or_create_access_tokens(self.entity_id) @callback def _async_update_and_write_token(self, now: datetime) -> None: _LOGGER.debug("Updating access tokens for %s", self.entity_id) self.async_update_token() self.async_write_ha_state() @callback def async_update_token(self) -> None: """Update the used token.""" self.access_tokens.append( hashlib.sha256(_RND.getrandbits(256).to_bytes(32, "little")).hexdigest() ) @callback def async_cleanup_tokens(self) -> None: """Clean up any remaining tokens on removal.""" if self.entity_id in self.data.access_tokens: del self.data.access_tokens[self.entity_id] async def async_added_to_hass(self) -> None: """Run when entity about to be added to hass. Injects callbacks to update access tokens automatically """ await super().async_added_to_hass() self.async_update_token() self.async_on_remove( self.hass.helpers.event.async_track_time_interval( self._async_update_and_write_token, TOKEN_CHANGE_INTERVAL ) ) self.async_on_remove(self.async_cleanup_tokens) class EventThumbnailMixin(AccessTokenMixin): """Adds motion event attributes to sensor.""" def __init__(self, *args: Any, **kwarg: Any) -> None: """Init an sensor that has event thumbnails.""" super().__init__(*args, **kwarg) self._event: Event | None = None @callback def _async_get_event(self) -> Event | None: """Get event from Protect device. To be overridden by child classes. """ raise NotImplementedError() @callback def _async_thumbnail_extra_attrs(self) -> dict[str, Any]: # Camera motion sensors with object detection attrs: dict[str, Any] = { ATTR_EVENT_SCORE: 0, ATTR_EVENT_THUMB: None, } if self._event is None: return attrs attrs[ATTR_EVENT_SCORE] = self._event.score if len(self.access_tokens) > 0: params = urlencode( {"entity_id": self.entity_id, "token": self.access_tokens[-1]} ) attrs[ATTR_EVENT_THUMB] = ( ThumbnailProxyView.url.format(event_id=self._event.id) + f"?{params}" ) return attrs @callback def _async_update_device_from_protect(self) -> None: super()._async_update_device_from_protect() self._event = self._async_get_event() attrs = self.extra_state_attributes or {} self._attr_extra_state_attributes = { **attrs, **self._async_thumbnail_extra_attrs(), }