114 lines
		
	
	
		
			3.7 KiB
		
	
	
	
		
			Python
		
	
	
			
		
		
	
	
			114 lines
		
	
	
		
			3.7 KiB
		
	
	
	
		
			Python
		
	
	
"""Image platform for UniFi Network integration.
 | 
						|
 | 
						|
Support for QR code for guest WLANs.
 | 
						|
"""
 | 
						|
 | 
						|
from __future__ import annotations
 | 
						|
 | 
						|
from collections.abc import Callable
 | 
						|
from dataclasses import dataclass
 | 
						|
 | 
						|
from aiounifi.interfaces.api_handlers import ItemEvent
 | 
						|
from aiounifi.interfaces.wlans import Wlans
 | 
						|
from aiounifi.models.api import ApiItemT
 | 
						|
from aiounifi.models.wlan import Wlan
 | 
						|
 | 
						|
from homeassistant.components.image import ImageEntity, ImageEntityDescription
 | 
						|
from homeassistant.const import EntityCategory
 | 
						|
from homeassistant.core import HomeAssistant, callback
 | 
						|
from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 | 
						|
from homeassistant.util import dt as dt_util
 | 
						|
 | 
						|
from . import UnifiConfigEntry
 | 
						|
from .entity import (
 | 
						|
    HandlerT,
 | 
						|
    UnifiEntity,
 | 
						|
    UnifiEntityDescription,
 | 
						|
    async_wlan_available_fn,
 | 
						|
    async_wlan_device_info_fn,
 | 
						|
)
 | 
						|
from .hub import UnifiHub
 | 
						|
 | 
						|
 | 
						|
@callback
 | 
						|
def async_wlan_qr_code_image_fn(hub: UnifiHub, wlan: Wlan) -> bytes:
 | 
						|
    """Calculate receiving data transfer value."""
 | 
						|
    return hub.api.wlans.generate_wlan_qr_code(wlan)
 | 
						|
 | 
						|
 | 
						|
@dataclass(frozen=True, kw_only=True)
 | 
						|
class UnifiImageEntityDescription(
 | 
						|
    ImageEntityDescription, UnifiEntityDescription[HandlerT, ApiItemT]
 | 
						|
):
 | 
						|
    """Class describing UniFi image entity."""
 | 
						|
 | 
						|
    image_fn: Callable[[UnifiHub, ApiItemT], bytes]
 | 
						|
    value_fn: Callable[[ApiItemT], str | None]
 | 
						|
 | 
						|
 | 
						|
ENTITY_DESCRIPTIONS: tuple[UnifiImageEntityDescription, ...] = (
 | 
						|
    UnifiImageEntityDescription[Wlans, Wlan](
 | 
						|
        key="WLAN QR Code",
 | 
						|
        translation_key="wlan_qr_code",
 | 
						|
        entity_category=EntityCategory.DIAGNOSTIC,
 | 
						|
        entity_registry_enabled_default=False,
 | 
						|
        api_handler_fn=lambda api: api.wlans,
 | 
						|
        available_fn=async_wlan_available_fn,
 | 
						|
        device_info_fn=async_wlan_device_info_fn,
 | 
						|
        name_fn=lambda wlan: "QR Code",
 | 
						|
        object_fn=lambda api, obj_id: api.wlans[obj_id],
 | 
						|
        unique_id_fn=lambda hub, obj_id: f"qr_code-{obj_id}",
 | 
						|
        image_fn=async_wlan_qr_code_image_fn,
 | 
						|
        value_fn=lambda obj: obj.x_passphrase,
 | 
						|
    ),
 | 
						|
)
 | 
						|
 | 
						|
 | 
						|
async def async_setup_entry(
 | 
						|
    hass: HomeAssistant,
 | 
						|
    config_entry: UnifiConfigEntry,
 | 
						|
    async_add_entities: AddConfigEntryEntitiesCallback,
 | 
						|
) -> None:
 | 
						|
    """Set up image platform for UniFi Network integration."""
 | 
						|
    config_entry.runtime_data.entity_loader.register_platform(
 | 
						|
        async_add_entities, UnifiImageEntity, ENTITY_DESCRIPTIONS, requires_admin=True
 | 
						|
    )
 | 
						|
 | 
						|
 | 
						|
class UnifiImageEntity(UnifiEntity[HandlerT, ApiItemT], ImageEntity):
 | 
						|
    """Base representation of a UniFi image."""
 | 
						|
 | 
						|
    entity_description: UnifiImageEntityDescription[HandlerT, ApiItemT]
 | 
						|
    _attr_content_type = "image/png"
 | 
						|
 | 
						|
    current_image: bytes | None = None
 | 
						|
    previous_value: str | None = None
 | 
						|
 | 
						|
    def __init__(
 | 
						|
        self,
 | 
						|
        obj_id: str,
 | 
						|
        hub: UnifiHub,
 | 
						|
        description: UnifiEntityDescription[HandlerT, ApiItemT],
 | 
						|
    ) -> None:
 | 
						|
        """Initiatlize UniFi Image entity."""
 | 
						|
        super().__init__(obj_id, hub, description)
 | 
						|
        ImageEntity.__init__(self, hub.hass)
 | 
						|
 | 
						|
    def image(self) -> bytes | None:
 | 
						|
        """Return bytes of image."""
 | 
						|
        if self.current_image is None:
 | 
						|
            description = self.entity_description
 | 
						|
            obj = description.object_fn(self.api, self._obj_id)
 | 
						|
            self.current_image = description.image_fn(self.hub, obj)
 | 
						|
        return self.current_image
 | 
						|
 | 
						|
    @callback
 | 
						|
    def async_update_state(self, event: ItemEvent, obj_id: str) -> None:
 | 
						|
        """Update entity state."""
 | 
						|
        description = self.entity_description
 | 
						|
        obj = description.object_fn(self.api, self._obj_id)
 | 
						|
        if (value := description.value_fn(obj)) != self.previous_value:
 | 
						|
            self.previous_value = value
 | 
						|
            self.current_image = None
 | 
						|
            self._attr_image_last_updated = dt_util.utcnow()
 |