"""Shared Entity definition for UniFi Protect Integration.""" from __future__ import annotations from collections.abc import Callable, Sequence import logging from typing import TYPE_CHECKING, Any from pyunifiprotect.data import ( NVR, Camera, Chime, Doorlock, Event, Light, ModelType, ProtectAdoptableDeviceModel, ProtectModelWithId, Sensor, StateType, Viewer, ) from homeassistant.core import callback import homeassistant.helpers.device_registry as dr from homeassistant.helpers.device_registry import DeviceInfo from homeassistant.helpers.entity import Entity, EntityDescription from homeassistant.helpers.typing import UNDEFINED from .const import ( ATTR_EVENT_ID, ATTR_EVENT_SCORE, DEFAULT_ATTRIBUTION, DEFAULT_BRAND, DOMAIN, ) from .data import ProtectData from .models import PermRequired, ProtectEventMixin, ProtectRequiredKeysMixin _LOGGER = logging.getLogger(__name__) @callback def _async_device_entities( data: ProtectData, klass: type[ProtectDeviceEntity], model_type: ModelType, descs: Sequence[ProtectRequiredKeysMixin], unadopted_descs: Sequence[ProtectRequiredKeysMixin], ufp_device: ProtectAdoptableDeviceModel | None = None, ) -> list[ProtectDeviceEntity]: if not descs and not unadopted_descs: return [] entities: list[ProtectDeviceEntity] = [] devices = ( [ufp_device] if ufp_device is not None else data.get_by_types({model_type}, ignore_unadopted=False) ) for device in devices: if TYPE_CHECKING: assert isinstance(device, (Camera, Light, Sensor, Viewer, Doorlock, Chime)) if not device.is_adopted_by_us: for description in unadopted_descs: entities.append( klass( data, device=device, description=description, ) ) _LOGGER.debug( "Adding %s entity %s for %s", klass.__name__, description.name, device.display_name, ) continue can_write = device.can_write(data.api.bootstrap.auth_user) for description in descs: if description.ufp_perm is not None: if description.ufp_perm is PermRequired.WRITE and not can_write: continue if description.ufp_perm is PermRequired.NO_WRITE and can_write: continue if ( description.ufp_perm is PermRequired.DELETE and not device.can_delete(data.api.bootstrap.auth_user) ): continue if not description.has_required(device): continue entities.append( klass( data, device=device, description=description, ) ) _LOGGER.debug( "Adding %s entity %s for %s", klass.__name__, description.name, device.display_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, lock_descs: Sequence[ProtectRequiredKeysMixin] | None = None, chime_descs: Sequence[ProtectRequiredKeysMixin] | None = None, all_descs: Sequence[ProtectRequiredKeysMixin] | None = None, unadopted_descs: Sequence[ProtectRequiredKeysMixin] | None = None, ufp_device: ProtectAdoptableDeviceModel | None = None, ) -> list[ProtectDeviceEntity]: """Generate a list of all the device entities.""" all_descs = list(all_descs or []) unadopted_descs = list(unadopted_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 lock_descs = list(lock_descs or []) + all_descs chime_descs = list(chime_descs or []) + all_descs if ufp_device is None: return ( _async_device_entities( data, klass, ModelType.CAMERA, camera_descs, unadopted_descs ) + _async_device_entities( data, klass, ModelType.LIGHT, light_descs, unadopted_descs ) + _async_device_entities( data, klass, ModelType.SENSOR, sense_descs, unadopted_descs ) + _async_device_entities( data, klass, ModelType.VIEWPORT, viewer_descs, unadopted_descs ) + _async_device_entities( data, klass, ModelType.DOORLOCK, lock_descs, unadopted_descs ) + _async_device_entities( data, klass, ModelType.CHIME, chime_descs, unadopted_descs ) ) descs = [] if ufp_device.model is ModelType.CAMERA: descs = camera_descs elif ufp_device.model is ModelType.LIGHT: descs = light_descs elif ufp_device.model is ModelType.SENSOR: descs = sense_descs elif ufp_device.model is ModelType.VIEWPORT: descs = viewer_descs elif ufp_device.model is ModelType.DOORLOCK: descs = lock_descs elif ufp_device.model is ModelType.CHIME: descs = chime_descs if not descs and not unadopted_descs or ufp_device.model is None: return [] return _async_device_entities( data, klass, ufp_device.model, descs, unadopted_descs, ufp_device ) 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 self._async_get_ufp_enabled: Callable[ [ProtectAdoptableDeviceModel], bool ] | None = None if description is None: self._attr_unique_id = f"{self.device.mac}" self._attr_name = f"{self.device.display_name}" else: self.entity_description = description self._attr_unique_id = f"{self.device.mac}_{description.key}" name = ( description.name if description.name and description.name is not UNDEFINED else "" ) self._attr_name = f"{self.device.display_name} {name.title()}" if isinstance(description, ProtectRequiredKeysMixin): self._async_get_ufp_enabled = description.get_ufp_enabled self._attr_attribution = DEFAULT_ATTRIBUTION self._async_set_device_info() self._async_update_device_from_protect(device) 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.display_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, device: ProtectModelWithId) -> None: """Update Entity object from Protect device.""" if TYPE_CHECKING: assert isinstance(device, ProtectAdoptableDeviceModel) if last_update_success := self.data.last_update_success: self.device = device async_get_ufp_enabled = self._async_get_ufp_enabled self._attr_available = ( last_update_success and ( device.state is StateType.CONNECTED or (not device.is_adopted_by_us and device.can_adopt) ) and (not async_get_ufp_enabled or async_get_ufp_enabled(device)) ) @callback def _async_get_state_attrs(self) -> tuple[Any, ...]: """Retrieve data that goes into the current state of the entity. Called before and after updating entity and state is only written if there is a change. """ return (self._attr_available,) @callback def _async_updated_event(self, device: ProtectModelWithId) -> None: """When device is updated from Protect.""" previous_attrs = self._async_get_state_attrs() self._async_update_device_from_protect(device) current_attrs = self._async_get_state_attrs() if previous_attrs != current_attrs: if _LOGGER.isEnabledFor(logging.DEBUG): device_name = device.name if hasattr(self, "entity_description") and self.entity_description.name: device_name += f" {self.entity_description.name}" _LOGGER.debug( "Updating state [%s (%s)] %s -> %s", device_name, device.mac, previous_attrs, current_attrs, ) 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.mac, self._async_updated_event ) ) class ProtectNVREntity(ProtectDeviceEntity): """Base class for unifi protect entities.""" # separate subclass on purpose device: NVR def __init__( self, entry: ProtectData, device: NVR, description: EntityDescription | None = None, ) -> None: """Initialize the entity.""" super().__init__(entry, device, description) @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.display_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, device: ProtectModelWithId) -> None: data = self.data last_update_success = data.last_update_success if last_update_success: self.device = data.api.bootstrap.nvr self._attr_available = last_update_success class EventEntityMixin(ProtectDeviceEntity): """Adds motion event attributes to sensor.""" _unrecorded_attributes = frozenset({ATTR_EVENT_ID, ATTR_EVENT_SCORE}) entity_description: ProtectEventMixin 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_update_device_from_protect(self, device: ProtectModelWithId) -> None: event = self.entity_description.get_event_obj(device) if event is not None: self._attr_extra_state_attributes = { ATTR_EVENT_ID: event.id, ATTR_EVENT_SCORE: event.score, } else: self._attr_extra_state_attributes = {} self._event = event super()._async_update_device_from_protect(device)