"""UniFi Controller abstraction.""" import asyncio import ssl import async_timeout from aiohttp import CookieJar import aiounifi from homeassistant.exceptions import ConfigEntryNotReady from homeassistant.const import CONF_HOST from homeassistant.helpers import aiohttp_client from homeassistant.helpers.dispatcher import async_dispatcher_send from .const import ( CONF_CONTROLLER, CONF_SITE_ID, CONTROLLER_ID, LOGGER, UNIFI_CONFIG) from .errors import AuthenticationRequired, CannotConnect 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._site_name = None self._site_role = None self.unifi_config = {} @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 @property def event_update(self): """Event specific per UniFi entry to signal new data.""" return 'unifi-update-{}'.format( CONTROLLER_ID.format(host=self.host, site=self.site)) async def request_update(self): """Request an update.""" if self.progress is not None: return await self.progress self.progress = self.hass.async_create_task(self.async_update()) await self.progress self.progress = None async def async_update(self): """Update UniFi controller information.""" failed = False try: with async_timeout.timeout(10): await self.api.clients.update() await self.api.devices.update() except aiounifi.LoginRequired: try: with async_timeout.timeout(5): await self.api.login() except (asyncio.TimeoutError, aiounifi.AiounifiException): failed = True if self.available: LOGGER.error('Unable to reach controller %s', self.host) self.available = False except (asyncio.TimeoutError, aiounifi.AiounifiException): failed = True if self.available: LOGGER.error('Unable to reach controller %s', self.host) self.available = False if not failed and not self.available: LOGGER.info('Reconnected to controller %s', self.host) self.available = True async_dispatcher_send(self.hass, self.event_update) async def async_setup(self): """Set up a UniFi controller.""" hass = self.hass try: self.api = await get_controller( self.hass, **self.config_entry.data[CONF_CONTROLLER]) await self.api.initialize() sites = await self.api.sites() for site in sites.values(): if self.site == site['name']: self._site_name = site['desc'] self._site_role = site['role'] break except CannotConnect: raise ConfigEntryNotReady except Exception: # pylint: disable=broad-except LOGGER.error( 'Unknown error connecting with UniFi controller.') return False for unifi_config in hass.data[UNIFI_CONFIG]: if self.host == unifi_config[CONF_HOST] and \ self.site == unifi_config[CONF_SITE_ID]: self.unifi_config = unifi_config break for platform in ['device_tracker', 'switch']: hass.async_create_task( hass.config_entries.async_forward_entry_setup( self.config_entry, platform)) return True async def async_reset(self): """Reset this controller to default state. Will cancel any scheduled setup retry and will unload the config entry. """ # If the authentication was wrong. if self.api is None: return True for platform in ['device_tracker', 'switch']: await self.hass.config_entries.async_forward_entry_unload( self.config_entry, platform) return True async def get_controller( hass, host, username, password, port, site, verify_ssl): """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 ) try: with async_timeout.timeout(10): await controller.login() return controller except aiounifi.Unauthorized: LOGGER.warning("Connected to UniFi at %s but not registered.", host) raise AuthenticationRequired except (asyncio.TimeoutError, aiounifi.RequestError): LOGGER.error("Error connecting to the UniFi controller at %s", host) raise CannotConnect except aiounifi.AiounifiException: LOGGER.exception('Unknown UniFi communication error occurred') raise AuthenticationRequired