* Remove unnecessary exception re-wraps * Preserve exception chains on re-raise We slap "from cause" to almost all possible cases here. In some cases it could conceivably be better to do "from None" if we really want to hide the cause. However those should be in the minority, and "from cause" should be an improvement over the corresponding raise without a "from" in all cases anyway. The only case where we raise from None here is in plex, where the exception for an original invalid SSL cert is not the root cause for failure to validate a newly fetched one. Follow local convention on exception variable names if there is a consistent one, otherwise `err` to match with majority of codebase. * Fix mistaken re-wrap in homematicip_cloud/hap.py Missed the difference between HmipConnectionError and HmipcConnectionError. * Do not hide original error on plex new cert validation error Original is not the cause for the new one, but showing old in the traceback is useful nevertheless.
441 lines
14 KiB
Python
441 lines
14 KiB
Python
"""UniFi Controller abstraction."""
|
|
import asyncio
|
|
from datetime import timedelta
|
|
import ssl
|
|
|
|
from aiohttp import CookieJar
|
|
import aiounifi
|
|
from aiounifi.controller import (
|
|
DATA_CLIENT_REMOVED,
|
|
DATA_EVENT,
|
|
SIGNAL_CONNECTION_STATE,
|
|
SIGNAL_DATA,
|
|
)
|
|
from aiounifi.events import (
|
|
ACCESS_POINT_CONNECTED,
|
|
GATEWAY_CONNECTED,
|
|
SWITCH_CONNECTED,
|
|
WIRED_CLIENT_CONNECTED,
|
|
WIRELESS_CLIENT_CONNECTED,
|
|
WIRELESS_GUEST_CONNECTED,
|
|
)
|
|
from aiounifi.websocket import STATE_DISCONNECTED, STATE_RUNNING
|
|
import async_timeout
|
|
|
|
from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
|
|
from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
|
|
from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
|
|
from homeassistant.const import CONF_HOST
|
|
from homeassistant.core import callback
|
|
from homeassistant.exceptions import ConfigEntryNotReady
|
|
from homeassistant.helpers import aiohttp_client
|
|
from homeassistant.helpers.dispatcher import async_dispatcher_send
|
|
|
|
from .const import (
|
|
CONF_ALLOW_BANDWIDTH_SENSORS,
|
|
CONF_BLOCK_CLIENT,
|
|
CONF_CONTROLLER,
|
|
CONF_DETECTION_TIME,
|
|
CONF_IGNORE_WIRED_BUG,
|
|
CONF_POE_CLIENTS,
|
|
CONF_SITE_ID,
|
|
CONF_SSID_FILTER,
|
|
CONF_TRACK_CLIENTS,
|
|
CONF_TRACK_DEVICES,
|
|
CONF_TRACK_WIRED_CLIENTS,
|
|
CONTROLLER_ID,
|
|
DEFAULT_ALLOW_BANDWIDTH_SENSORS,
|
|
DEFAULT_DETECTION_TIME,
|
|
DEFAULT_IGNORE_WIRED_BUG,
|
|
DEFAULT_POE_CLIENTS,
|
|
DEFAULT_TRACK_CLIENTS,
|
|
DEFAULT_TRACK_DEVICES,
|
|
DEFAULT_TRACK_WIRED_CLIENTS,
|
|
DOMAIN as UNIFI_DOMAIN,
|
|
LOGGER,
|
|
UNIFI_WIRELESS_CLIENTS,
|
|
)
|
|
from .errors import AuthenticationRequired, CannotConnect
|
|
|
|
RETRY_TIMER = 15
|
|
SUPPORTED_PLATFORMS = [TRACKER_DOMAIN, SENSOR_DOMAIN, SWITCH_DOMAIN]
|
|
|
|
CLIENT_CONNECTED = (
|
|
WIRED_CLIENT_CONNECTED,
|
|
WIRELESS_CLIENT_CONNECTED,
|
|
WIRELESS_GUEST_CONNECTED,
|
|
)
|
|
DEVICE_CONNECTED = (
|
|
ACCESS_POINT_CONNECTED,
|
|
GATEWAY_CONNECTED,
|
|
SWITCH_CONNECTED,
|
|
)
|
|
|
|
|
|
class UniFiController:
|
|
"""Manages a single UniFi Controller."""
|
|
|
|
def __init__(self, hass, config_entry):
|
|
"""Initialize the system."""
|
|
self.hass = hass
|
|
self.config_entry = config_entry
|
|
self.available = True
|
|
self.api = None
|
|
self.progress = None
|
|
self.wireless_clients = None
|
|
|
|
self.listeners = []
|
|
self._site_name = None
|
|
self._site_role = None
|
|
|
|
self.entities = {}
|
|
|
|
@property
|
|
def controller_id(self):
|
|
"""Return the controller ID."""
|
|
return CONTROLLER_ID.format(host=self.host, site=self.site)
|
|
|
|
@property
|
|
def host(self):
|
|
"""Return the host of this controller."""
|
|
return self.config_entry.data[CONF_CONTROLLER][CONF_HOST]
|
|
|
|
@property
|
|
def site(self):
|
|
"""Return the site of this config entry."""
|
|
return self.config_entry.data[CONF_CONTROLLER][CONF_SITE_ID]
|
|
|
|
@property
|
|
def site_name(self):
|
|
"""Return the nice name of site."""
|
|
return self._site_name
|
|
|
|
@property
|
|
def site_role(self):
|
|
"""Return the site user role of this controller."""
|
|
return self._site_role
|
|
|
|
@property
|
|
def mac(self):
|
|
"""Return the mac address of this controller."""
|
|
for client in self.api.clients.values():
|
|
if self.host == client.ip:
|
|
return client.mac
|
|
return None
|
|
|
|
# Device tracker options
|
|
|
|
@property
|
|
def option_track_clients(self):
|
|
"""Config entry option to not track clients."""
|
|
return self.config_entry.options.get(CONF_TRACK_CLIENTS, DEFAULT_TRACK_CLIENTS)
|
|
|
|
@property
|
|
def option_track_wired_clients(self):
|
|
"""Config entry option to not track wired clients."""
|
|
return self.config_entry.options.get(
|
|
CONF_TRACK_WIRED_CLIENTS, DEFAULT_TRACK_WIRED_CLIENTS
|
|
)
|
|
|
|
@property
|
|
def option_track_devices(self):
|
|
"""Config entry option to not track devices."""
|
|
return self.config_entry.options.get(CONF_TRACK_DEVICES, DEFAULT_TRACK_DEVICES)
|
|
|
|
@property
|
|
def option_ssid_filter(self):
|
|
"""Config entry option listing what SSIDs are being used to track clients."""
|
|
return self.config_entry.options.get(CONF_SSID_FILTER, [])
|
|
|
|
@property
|
|
def option_detection_time(self):
|
|
"""Config entry option defining number of seconds from last seen to away."""
|
|
return timedelta(
|
|
seconds=self.config_entry.options.get(
|
|
CONF_DETECTION_TIME, DEFAULT_DETECTION_TIME
|
|
)
|
|
)
|
|
|
|
@property
|
|
def option_ignore_wired_bug(self):
|
|
"""Config entry option to ignore wired bug."""
|
|
return self.config_entry.options.get(
|
|
CONF_IGNORE_WIRED_BUG, DEFAULT_IGNORE_WIRED_BUG
|
|
)
|
|
|
|
# Client control options
|
|
|
|
@property
|
|
def option_poe_clients(self):
|
|
"""Config entry option to control poe clients."""
|
|
return self.config_entry.options.get(CONF_POE_CLIENTS, DEFAULT_POE_CLIENTS)
|
|
|
|
@property
|
|
def option_block_clients(self):
|
|
"""Config entry option with list of clients to control network access."""
|
|
return self.config_entry.options.get(CONF_BLOCK_CLIENT, [])
|
|
|
|
# Statistics sensor options
|
|
|
|
@property
|
|
def option_allow_bandwidth_sensors(self):
|
|
"""Config entry option to allow bandwidth sensors."""
|
|
return self.config_entry.options.get(
|
|
CONF_ALLOW_BANDWIDTH_SENSORS, DEFAULT_ALLOW_BANDWIDTH_SENSORS
|
|
)
|
|
|
|
@callback
|
|
def async_unifi_signalling_callback(self, signal, data):
|
|
"""Handle messages back from UniFi library."""
|
|
if signal == SIGNAL_CONNECTION_STATE:
|
|
|
|
if data == STATE_DISCONNECTED and self.available:
|
|
LOGGER.warning("Lost connection to UniFi controller")
|
|
|
|
if (data == STATE_RUNNING and not self.available) or (
|
|
data == STATE_DISCONNECTED and self.available
|
|
):
|
|
self.available = data == STATE_RUNNING
|
|
async_dispatcher_send(self.hass, self.signal_reachable)
|
|
|
|
if not self.available:
|
|
self.hass.loop.call_later(RETRY_TIMER, self.reconnect, True)
|
|
else:
|
|
LOGGER.info("Connected to UniFi controller")
|
|
|
|
elif signal == SIGNAL_DATA and data:
|
|
|
|
if DATA_EVENT in data:
|
|
clients_connected = set()
|
|
devices_connected = set()
|
|
wireless_clients_connected = False
|
|
|
|
for event in data[DATA_EVENT]:
|
|
|
|
if event.event in CLIENT_CONNECTED:
|
|
clients_connected.add(event.mac)
|
|
|
|
if not wireless_clients_connected and event.event in (
|
|
WIRELESS_CLIENT_CONNECTED,
|
|
WIRELESS_GUEST_CONNECTED,
|
|
):
|
|
wireless_clients_connected = True
|
|
|
|
elif event.event in DEVICE_CONNECTED:
|
|
devices_connected.add(event.mac)
|
|
|
|
if wireless_clients_connected:
|
|
self.update_wireless_clients()
|
|
if clients_connected or devices_connected:
|
|
async_dispatcher_send(
|
|
self.hass,
|
|
self.signal_update,
|
|
clients_connected,
|
|
devices_connected,
|
|
)
|
|
|
|
elif DATA_CLIENT_REMOVED in data:
|
|
async_dispatcher_send(
|
|
self.hass, self.signal_remove, data[DATA_CLIENT_REMOVED]
|
|
)
|
|
|
|
@property
|
|
def signal_reachable(self) -> str:
|
|
"""Integration specific event to signal a change in connection status."""
|
|
return f"unifi-reachable-{self.controller_id}"
|
|
|
|
@property
|
|
def signal_update(self):
|
|
"""Event specific per UniFi entry to signal new data."""
|
|
return f"unifi-update-{self.controller_id}"
|
|
|
|
@property
|
|
def signal_remove(self):
|
|
"""Event specific per UniFi entry to signal removal of entities."""
|
|
return f"unifi-remove-{self.controller_id}"
|
|
|
|
@property
|
|
def signal_options_update(self):
|
|
"""Event specific per UniFi entry to signal new options."""
|
|
return f"unifi-options-{self.controller_id}"
|
|
|
|
def update_wireless_clients(self):
|
|
"""Update set of known to be wireless clients."""
|
|
new_wireless_clients = set()
|
|
|
|
for client_id in self.api.clients:
|
|
if (
|
|
client_id not in self.wireless_clients
|
|
and not self.api.clients[client_id].is_wired
|
|
):
|
|
new_wireless_clients.add(client_id)
|
|
|
|
if new_wireless_clients:
|
|
self.wireless_clients |= new_wireless_clients
|
|
unifi_wireless_clients = self.hass.data[UNIFI_WIRELESS_CLIENTS]
|
|
unifi_wireless_clients.update_data(self.wireless_clients, self.config_entry)
|
|
|
|
async def async_setup(self):
|
|
"""Set up a UniFi controller."""
|
|
try:
|
|
self.api = await get_controller(
|
|
self.hass,
|
|
**self.config_entry.data[CONF_CONTROLLER],
|
|
async_callback=self.async_unifi_signalling_callback,
|
|
)
|
|
await self.api.initialize()
|
|
|
|
sites = await self.api.sites()
|
|
|
|
for site in sites.values():
|
|
if self.site == site["name"]:
|
|
self._site_name = site["desc"]
|
|
break
|
|
|
|
description = await self.api.site_description()
|
|
self._site_role = description[0]["site_role"]
|
|
|
|
except CannotConnect as err:
|
|
raise ConfigEntryNotReady from err
|
|
|
|
except Exception as err: # pylint: disable=broad-except
|
|
LOGGER.error("Unknown error connecting with UniFi controller: %s", err)
|
|
return False
|
|
|
|
# Restore clients that is not a part of active clients list.
|
|
entity_registry = await self.hass.helpers.entity_registry.async_get_registry()
|
|
for entity in entity_registry.entities.values():
|
|
if (
|
|
entity.config_entry_id != self.config_entry.entry_id
|
|
or "-" not in entity.unique_id
|
|
):
|
|
continue
|
|
|
|
mac = ""
|
|
if entity.domain == TRACKER_DOMAIN:
|
|
mac, _ = entity.unique_id.split("-", 1)
|
|
elif entity.domain == SWITCH_DOMAIN:
|
|
_, mac = entity.unique_id.split("-", 1)
|
|
|
|
if mac in self.api.clients or mac not in self.api.clients_all:
|
|
continue
|
|
|
|
client = self.api.clients_all[mac]
|
|
self.api.clients.process_raw([client.raw])
|
|
LOGGER.debug(
|
|
"Restore disconnected client %s (%s)",
|
|
entity.entity_id,
|
|
client.mac,
|
|
)
|
|
|
|
wireless_clients = self.hass.data[UNIFI_WIRELESS_CLIENTS]
|
|
self.wireless_clients = wireless_clients.get_data(self.config_entry)
|
|
self.update_wireless_clients()
|
|
|
|
for platform in SUPPORTED_PLATFORMS:
|
|
self.hass.async_create_task(
|
|
self.hass.config_entries.async_forward_entry_setup(
|
|
self.config_entry, platform
|
|
)
|
|
)
|
|
|
|
self.api.start_websocket()
|
|
|
|
self.config_entry.add_update_listener(self.async_config_entry_updated)
|
|
|
|
return True
|
|
|
|
@staticmethod
|
|
async def async_config_entry_updated(hass, config_entry) -> None:
|
|
"""Handle signals of config entry being updated."""
|
|
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
|
|
async_dispatcher_send(hass, controller.signal_options_update)
|
|
|
|
@callback
|
|
def reconnect(self, log=False) -> None:
|
|
"""Prepare to reconnect UniFi session."""
|
|
if log:
|
|
LOGGER.info("Will try to reconnect to UniFi controller")
|
|
self.hass.loop.create_task(self.async_reconnect())
|
|
|
|
async def async_reconnect(self) -> None:
|
|
"""Try to reconnect UniFi session."""
|
|
try:
|
|
with async_timeout.timeout(5):
|
|
await self.api.login()
|
|
self.api.start_websocket()
|
|
|
|
except (asyncio.TimeoutError, aiounifi.AiounifiException):
|
|
self.hass.loop.call_later(RETRY_TIMER, self.reconnect)
|
|
|
|
@callback
|
|
def shutdown(self, event) -> None:
|
|
"""Wrap the call to unifi.close.
|
|
|
|
Used as an argument to EventBus.async_listen_once.
|
|
"""
|
|
self.api.stop_websocket()
|
|
|
|
async def async_reset(self):
|
|
"""Reset this controller to default state.
|
|
|
|
Will cancel any scheduled setup retry and will unload
|
|
the config entry.
|
|
"""
|
|
self.api.stop_websocket()
|
|
|
|
for platform in SUPPORTED_PLATFORMS:
|
|
await self.hass.config_entries.async_forward_entry_unload(
|
|
self.config_entry, platform
|
|
)
|
|
|
|
for unsub_dispatcher in self.listeners:
|
|
unsub_dispatcher()
|
|
self.listeners = []
|
|
|
|
return True
|
|
|
|
|
|
async def get_controller(
|
|
hass, host, username, password, port, site, verify_ssl, async_callback=None
|
|
):
|
|
"""Create a controller object and verify authentication."""
|
|
sslcontext = None
|
|
|
|
if verify_ssl:
|
|
session = aiohttp_client.async_get_clientsession(hass)
|
|
if isinstance(verify_ssl, str):
|
|
sslcontext = ssl.create_default_context(cafile=verify_ssl)
|
|
else:
|
|
session = aiohttp_client.async_create_clientsession(
|
|
hass, verify_ssl=verify_ssl, cookie_jar=CookieJar(unsafe=True)
|
|
)
|
|
|
|
controller = aiounifi.Controller(
|
|
host,
|
|
username=username,
|
|
password=password,
|
|
port=port,
|
|
site=site,
|
|
websession=session,
|
|
sslcontext=sslcontext,
|
|
callback=async_callback,
|
|
)
|
|
|
|
try:
|
|
with async_timeout.timeout(10):
|
|
await controller.check_unifi_os()
|
|
await controller.login()
|
|
return controller
|
|
|
|
except aiounifi.Unauthorized as err:
|
|
LOGGER.warning("Connected to UniFi at %s but not registered.", host)
|
|
raise AuthenticationRequired from err
|
|
|
|
except (asyncio.TimeoutError, aiounifi.RequestError) as err:
|
|
LOGGER.error("Error connecting to the UniFi controller at %s", host)
|
|
raise CannotConnect from err
|
|
|
|
except aiounifi.AiounifiException as err:
|
|
LOGGER.exception("Unknown UniFi communication error occurred")
|
|
raise AuthenticationRequired from err
|