mirror of
https://github.com/home-assistant/core
synced 2024-08-28 03:36:46 +02:00
432 lines
16 KiB
Python
432 lines
16 KiB
Python
"""Switch platform for UniFi Network integration.
|
|
|
|
Support for controlling power supply of clients which are powered over Ethernet (POE).
|
|
Support for controlling network access of clients selected in option flow.
|
|
Support for controlling deep packet inspection (DPI) restriction groups.
|
|
Support for controlling WLAN availability.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Callable, Coroutine
|
|
from dataclasses import dataclass
|
|
from typing import Any, Generic
|
|
|
|
import aiounifi
|
|
from aiounifi.interfaces.api_handlers import ItemEvent
|
|
from aiounifi.interfaces.clients import Clients
|
|
from aiounifi.interfaces.dpi_restriction_groups import DPIRestrictionGroups
|
|
from aiounifi.interfaces.outlets import Outlets
|
|
from aiounifi.interfaces.port_forwarding import PortForwarding
|
|
from aiounifi.interfaces.ports import Ports
|
|
from aiounifi.interfaces.wlans import Wlans
|
|
from aiounifi.models.api import ApiItemT
|
|
from aiounifi.models.client import Client, ClientBlockRequest
|
|
from aiounifi.models.device import DeviceSetOutletRelayRequest
|
|
from aiounifi.models.dpi_restriction_app import DPIRestrictionAppEnableRequest
|
|
from aiounifi.models.dpi_restriction_group import DPIRestrictionGroup
|
|
from aiounifi.models.event import Event, EventKey
|
|
from aiounifi.models.outlet import Outlet
|
|
from aiounifi.models.port import Port
|
|
from aiounifi.models.port_forward import PortForward, PortForwardEnableRequest
|
|
from aiounifi.models.wlan import Wlan, WlanEnableRequest
|
|
|
|
from homeassistant.components.switch import (
|
|
DOMAIN,
|
|
SwitchDeviceClass,
|
|
SwitchEntity,
|
|
SwitchEntityDescription,
|
|
)
|
|
from homeassistant.config_entries import ConfigEntry
|
|
from homeassistant.const import EntityCategory
|
|
from homeassistant.core import HomeAssistant, callback
|
|
from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
|
|
from homeassistant.helpers.entity_platform import AddEntitiesCallback
|
|
import homeassistant.helpers.entity_registry as er
|
|
|
|
from .const import ATTR_MANUFACTURER
|
|
from .controller import UNIFI_DOMAIN, UniFiController
|
|
from .entity import (
|
|
HandlerT,
|
|
SubscriptionT,
|
|
UnifiEntity,
|
|
UnifiEntityDescription,
|
|
async_client_device_info_fn,
|
|
async_device_available_fn,
|
|
async_device_device_info_fn,
|
|
async_wlan_device_info_fn,
|
|
)
|
|
|
|
CLIENT_BLOCKED = (EventKey.WIRED_CLIENT_BLOCKED, EventKey.WIRELESS_CLIENT_BLOCKED)
|
|
CLIENT_UNBLOCKED = (EventKey.WIRED_CLIENT_UNBLOCKED, EventKey.WIRELESS_CLIENT_UNBLOCKED)
|
|
|
|
|
|
@callback
|
|
def async_block_client_allowed_fn(controller: UniFiController, obj_id: str) -> bool:
|
|
"""Check if client is allowed."""
|
|
if obj_id in controller.option_supported_clients:
|
|
return True
|
|
return obj_id in controller.option_block_clients
|
|
|
|
|
|
@callback
|
|
def async_dpi_group_is_on_fn(
|
|
controller: UniFiController, dpi_group: DPIRestrictionGroup
|
|
) -> bool:
|
|
"""Calculate if all apps are enabled."""
|
|
api = controller.api
|
|
return all(
|
|
api.dpi_apps[app_id].enabled
|
|
for app_id in dpi_group.dpiapp_ids or []
|
|
if app_id in api.dpi_apps
|
|
)
|
|
|
|
|
|
@callback
|
|
def async_dpi_group_device_info_fn(
|
|
controller: UniFiController, obj_id: str
|
|
) -> DeviceInfo:
|
|
"""Create device registry entry for DPI group."""
|
|
return DeviceInfo(
|
|
entry_type=DeviceEntryType.SERVICE,
|
|
identifiers={(DOMAIN, f"unifi_controller_{obj_id}")},
|
|
manufacturer=ATTR_MANUFACTURER,
|
|
model="UniFi Network",
|
|
name="UniFi Network",
|
|
)
|
|
|
|
|
|
@callback
|
|
def async_port_forward_device_info_fn(
|
|
controller: UniFiController, obj_id: str
|
|
) -> DeviceInfo:
|
|
"""Create device registry entry for port forward."""
|
|
unique_id = controller.config_entry.unique_id
|
|
assert unique_id is not None
|
|
return DeviceInfo(
|
|
entry_type=DeviceEntryType.SERVICE,
|
|
identifiers={(DOMAIN, unique_id)},
|
|
manufacturer=ATTR_MANUFACTURER,
|
|
model="UniFi Network",
|
|
name="UniFi Network",
|
|
)
|
|
|
|
|
|
async def async_block_client_control_fn(
|
|
controller: UniFiController, obj_id: str, target: bool
|
|
) -> None:
|
|
"""Control network access of client."""
|
|
await controller.api.request(ClientBlockRequest.create(obj_id, not target))
|
|
|
|
|
|
async def async_dpi_group_control_fn(
|
|
controller: UniFiController, obj_id: str, target: bool
|
|
) -> None:
|
|
"""Enable or disable DPI group."""
|
|
dpi_group = controller.api.dpi_groups[obj_id]
|
|
await asyncio.gather(
|
|
*[
|
|
controller.api.request(
|
|
DPIRestrictionAppEnableRequest.create(app_id, target)
|
|
)
|
|
for app_id in dpi_group.dpiapp_ids or []
|
|
]
|
|
)
|
|
|
|
|
|
@callback
|
|
def async_outlet_supports_switching_fn(
|
|
controller: UniFiController, obj_id: str
|
|
) -> bool:
|
|
"""Determine if an outlet supports switching."""
|
|
outlet = controller.api.outlets[obj_id]
|
|
return outlet.has_relay or outlet.caps in (1, 3)
|
|
|
|
|
|
async def async_outlet_control_fn(
|
|
controller: UniFiController, obj_id: str, target: bool
|
|
) -> None:
|
|
"""Control outlet relay."""
|
|
mac, _, index = obj_id.partition("_")
|
|
device = controller.api.devices[mac]
|
|
await controller.api.request(
|
|
DeviceSetOutletRelayRequest.create(device, int(index), target)
|
|
)
|
|
|
|
|
|
async def async_poe_port_control_fn(
|
|
controller: UniFiController, obj_id: str, target: bool
|
|
) -> None:
|
|
"""Control poe state."""
|
|
mac, _, index = obj_id.partition("_")
|
|
port = controller.api.ports[obj_id]
|
|
on_state = "auto" if port.raw["poe_caps"] != 8 else "passthrough"
|
|
state = on_state if target else "off"
|
|
controller.async_queue_poe_port_command(mac, int(index), state)
|
|
|
|
|
|
async def async_port_forward_control_fn(
|
|
controller: UniFiController, obj_id: str, target: bool
|
|
) -> None:
|
|
"""Control port forward state."""
|
|
port_forward = controller.api.port_forwarding[obj_id]
|
|
await controller.api.request(PortForwardEnableRequest.create(port_forward, target))
|
|
|
|
|
|
async def async_wlan_control_fn(
|
|
controller: UniFiController, obj_id: str, target: bool
|
|
) -> None:
|
|
"""Control outlet relay."""
|
|
await controller.api.request(WlanEnableRequest.create(obj_id, target))
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class UnifiSwitchEntityDescriptionMixin(Generic[HandlerT, ApiItemT]):
|
|
"""Validate and load entities from different UniFi handlers."""
|
|
|
|
control_fn: Callable[[UniFiController, str, bool], Coroutine[Any, Any, None]]
|
|
is_on_fn: Callable[[UniFiController, ApiItemT], bool]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class UnifiSwitchEntityDescription(
|
|
SwitchEntityDescription,
|
|
UnifiEntityDescription[HandlerT, ApiItemT],
|
|
UnifiSwitchEntityDescriptionMixin[HandlerT, ApiItemT],
|
|
):
|
|
"""Class describing UniFi switch entity."""
|
|
|
|
custom_subscribe: Callable[[aiounifi.Controller], SubscriptionT] | None = None
|
|
only_event_for_state_change: bool = False
|
|
|
|
|
|
ENTITY_DESCRIPTIONS: tuple[UnifiSwitchEntityDescription, ...] = (
|
|
UnifiSwitchEntityDescription[Clients, Client](
|
|
key="Block client",
|
|
device_class=SwitchDeviceClass.SWITCH,
|
|
entity_category=EntityCategory.CONFIG,
|
|
has_entity_name=True,
|
|
icon="mdi:ethernet",
|
|
allowed_fn=async_block_client_allowed_fn,
|
|
api_handler_fn=lambda api: api.clients,
|
|
available_fn=lambda controller, obj_id: controller.available,
|
|
control_fn=async_block_client_control_fn,
|
|
device_info_fn=async_client_device_info_fn,
|
|
event_is_on=CLIENT_UNBLOCKED,
|
|
event_to_subscribe=CLIENT_BLOCKED + CLIENT_UNBLOCKED,
|
|
is_on_fn=lambda controller, client: not client.blocked,
|
|
name_fn=lambda client: None,
|
|
object_fn=lambda api, obj_id: api.clients[obj_id],
|
|
only_event_for_state_change=True,
|
|
should_poll=False,
|
|
supported_fn=lambda controller, obj_id: True,
|
|
unique_id_fn=lambda controller, obj_id: f"block-{obj_id}",
|
|
),
|
|
UnifiSwitchEntityDescription[DPIRestrictionGroups, DPIRestrictionGroup](
|
|
key="DPI restriction",
|
|
entity_category=EntityCategory.CONFIG,
|
|
icon="mdi:network",
|
|
allowed_fn=lambda controller, obj_id: controller.option_dpi_restrictions,
|
|
api_handler_fn=lambda api: api.dpi_groups,
|
|
available_fn=lambda controller, obj_id: controller.available,
|
|
control_fn=async_dpi_group_control_fn,
|
|
custom_subscribe=lambda api: api.dpi_apps.subscribe,
|
|
device_info_fn=async_dpi_group_device_info_fn,
|
|
event_is_on=None,
|
|
event_to_subscribe=None,
|
|
is_on_fn=async_dpi_group_is_on_fn,
|
|
name_fn=lambda group: group.name,
|
|
object_fn=lambda api, obj_id: api.dpi_groups[obj_id],
|
|
should_poll=False,
|
|
supported_fn=lambda c, obj_id: bool(c.api.dpi_groups[obj_id].dpiapp_ids),
|
|
unique_id_fn=lambda controller, obj_id: obj_id,
|
|
),
|
|
UnifiSwitchEntityDescription[Outlets, Outlet](
|
|
key="Outlet control",
|
|
device_class=SwitchDeviceClass.OUTLET,
|
|
has_entity_name=True,
|
|
allowed_fn=lambda controller, obj_id: True,
|
|
api_handler_fn=lambda api: api.outlets,
|
|
available_fn=async_device_available_fn,
|
|
control_fn=async_outlet_control_fn,
|
|
device_info_fn=async_device_device_info_fn,
|
|
event_is_on=None,
|
|
event_to_subscribe=None,
|
|
is_on_fn=lambda controller, outlet: outlet.relay_state,
|
|
name_fn=lambda outlet: outlet.name,
|
|
object_fn=lambda api, obj_id: api.outlets[obj_id],
|
|
should_poll=False,
|
|
supported_fn=async_outlet_supports_switching_fn,
|
|
unique_id_fn=lambda controller, obj_id: f"outlet-{obj_id}",
|
|
),
|
|
UnifiSwitchEntityDescription[PortForwarding, PortForward](
|
|
key="Port forward control",
|
|
device_class=SwitchDeviceClass.SWITCH,
|
|
entity_category=EntityCategory.CONFIG,
|
|
has_entity_name=True,
|
|
icon="mdi:upload-network",
|
|
allowed_fn=lambda controller, obj_id: True,
|
|
api_handler_fn=lambda api: api.port_forwarding,
|
|
available_fn=lambda controller, obj_id: controller.available,
|
|
control_fn=async_port_forward_control_fn,
|
|
device_info_fn=async_port_forward_device_info_fn,
|
|
event_is_on=None,
|
|
event_to_subscribe=None,
|
|
is_on_fn=lambda controller, port_forward: port_forward.enabled,
|
|
name_fn=lambda port_forward: f"{port_forward.name}",
|
|
object_fn=lambda api, obj_id: api.port_forwarding[obj_id],
|
|
should_poll=False,
|
|
supported_fn=lambda controller, obj_id: True,
|
|
unique_id_fn=lambda controller, obj_id: f"port_forward-{obj_id}",
|
|
),
|
|
UnifiSwitchEntityDescription[Ports, Port](
|
|
key="PoE port control",
|
|
device_class=SwitchDeviceClass.OUTLET,
|
|
entity_category=EntityCategory.CONFIG,
|
|
has_entity_name=True,
|
|
entity_registry_enabled_default=False,
|
|
icon="mdi:ethernet",
|
|
allowed_fn=lambda controller, obj_id: True,
|
|
api_handler_fn=lambda api: api.ports,
|
|
available_fn=async_device_available_fn,
|
|
control_fn=async_poe_port_control_fn,
|
|
device_info_fn=async_device_device_info_fn,
|
|
event_is_on=None,
|
|
event_to_subscribe=None,
|
|
is_on_fn=lambda controller, port: port.poe_mode != "off",
|
|
name_fn=lambda port: f"{port.name} PoE",
|
|
object_fn=lambda api, obj_id: api.ports[obj_id],
|
|
should_poll=False,
|
|
supported_fn=lambda controller, obj_id: controller.api.ports[obj_id].port_poe,
|
|
unique_id_fn=lambda controller, obj_id: f"poe-{obj_id}",
|
|
),
|
|
UnifiSwitchEntityDescription[Wlans, Wlan](
|
|
key="WLAN control",
|
|
device_class=SwitchDeviceClass.SWITCH,
|
|
entity_category=EntityCategory.CONFIG,
|
|
has_entity_name=True,
|
|
icon="mdi:wifi-check",
|
|
allowed_fn=lambda controller, obj_id: True,
|
|
api_handler_fn=lambda api: api.wlans,
|
|
available_fn=lambda controller, _: controller.available,
|
|
control_fn=async_wlan_control_fn,
|
|
device_info_fn=async_wlan_device_info_fn,
|
|
event_is_on=None,
|
|
event_to_subscribe=None,
|
|
is_on_fn=lambda controller, wlan: wlan.enabled,
|
|
name_fn=lambda wlan: None,
|
|
object_fn=lambda api, obj_id: api.wlans[obj_id],
|
|
should_poll=False,
|
|
supported_fn=lambda controller, obj_id: True,
|
|
unique_id_fn=lambda controller, obj_id: f"wlan-{obj_id}",
|
|
),
|
|
)
|
|
|
|
|
|
@callback
|
|
def async_update_unique_id(hass: HomeAssistant, config_entry: ConfigEntry) -> None:
|
|
"""Normalize switch unique ID to have a prefix rather than midfix.
|
|
|
|
Introduced with release 2023.12.
|
|
"""
|
|
controller: UniFiController = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
|
|
ent_reg = er.async_get(hass)
|
|
|
|
@callback
|
|
def update_unique_id(obj_id: str, type_name: str) -> None:
|
|
"""Rework unique ID."""
|
|
new_unique_id = f"{type_name}-{obj_id}"
|
|
if ent_reg.async_get_entity_id(DOMAIN, UNIFI_DOMAIN, new_unique_id):
|
|
return
|
|
|
|
prefix, _, suffix = obj_id.partition("_")
|
|
unique_id = f"{prefix}-{type_name}-{suffix}"
|
|
if entity_id := ent_reg.async_get_entity_id(DOMAIN, UNIFI_DOMAIN, unique_id):
|
|
ent_reg.async_update_entity(entity_id, new_unique_id=new_unique_id)
|
|
|
|
for obj_id in controller.api.outlets:
|
|
update_unique_id(obj_id, "outlet")
|
|
|
|
for obj_id in controller.api.ports:
|
|
update_unique_id(obj_id, "poe")
|
|
|
|
|
|
async def async_setup_entry(
|
|
hass: HomeAssistant,
|
|
config_entry: ConfigEntry,
|
|
async_add_entities: AddEntitiesCallback,
|
|
) -> None:
|
|
"""Set up switches for UniFi Network integration."""
|
|
async_update_unique_id(hass, config_entry)
|
|
UniFiController.register_platform(
|
|
hass,
|
|
config_entry,
|
|
async_add_entities,
|
|
UnifiSwitchEntity,
|
|
ENTITY_DESCRIPTIONS,
|
|
requires_admin=True,
|
|
)
|
|
|
|
|
|
class UnifiSwitchEntity(UnifiEntity[HandlerT, ApiItemT], SwitchEntity):
|
|
"""Base representation of a UniFi switch."""
|
|
|
|
entity_description: UnifiSwitchEntityDescription[HandlerT, ApiItemT]
|
|
only_event_for_state_change = False
|
|
|
|
@callback
|
|
def async_initiate_state(self) -> None:
|
|
"""Initiate entity state."""
|
|
self.async_update_state(ItemEvent.ADDED, self._obj_id)
|
|
self.only_event_for_state_change = (
|
|
self.entity_description.only_event_for_state_change
|
|
)
|
|
|
|
async def async_turn_on(self, **kwargs: Any) -> None:
|
|
"""Turn on switch."""
|
|
await self.entity_description.control_fn(self.controller, self._obj_id, True)
|
|
|
|
async def async_turn_off(self, **kwargs: Any) -> None:
|
|
"""Turn off switch."""
|
|
await self.entity_description.control_fn(self.controller, self._obj_id, False)
|
|
|
|
@callback
|
|
def async_update_state(self, event: ItemEvent, obj_id: str) -> None:
|
|
"""Update entity state.
|
|
|
|
Update attr_is_on.
|
|
"""
|
|
if self.only_event_for_state_change:
|
|
return
|
|
|
|
description = self.entity_description
|
|
obj = description.object_fn(self.controller.api, self._obj_id)
|
|
if (is_on := description.is_on_fn(self.controller, obj)) != self.is_on:
|
|
self._attr_is_on = is_on
|
|
|
|
@callback
|
|
def async_event_callback(self, event: Event) -> None:
|
|
"""Event subscription callback."""
|
|
if event.mac != self._obj_id:
|
|
return
|
|
|
|
description = self.entity_description
|
|
assert isinstance(description.event_to_subscribe, tuple)
|
|
assert isinstance(description.event_is_on, tuple)
|
|
|
|
if event.key in description.event_to_subscribe:
|
|
self._attr_is_on = event.key in description.event_is_on
|
|
self._attr_available = description.available_fn(self.controller, self._obj_id)
|
|
self.async_write_ha_state()
|
|
|
|
async def async_added_to_hass(self) -> None:
|
|
"""Register callbacks."""
|
|
await super().async_added_to_hass()
|
|
|
|
if self.entity_description.custom_subscribe is not None:
|
|
self.async_on_remove(
|
|
self.entity_description.custom_subscribe(self.controller.api)(
|
|
self.async_signalling_callback, ItemEvent.CHANGED
|
|
),
|
|
)
|