From b6934f0cd0c34273e902fb6ebf1adac91f816270 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Thu, 25 Jul 2019 16:56:56 +0200 Subject: [PATCH] UniFi block clients (#25478) * Allow blocking clients on UniFi networks --- homeassistant/components/unifi/__init__.py | 8 +- homeassistant/components/unifi/const.py | 1 + homeassistant/components/unifi/controller.py | 16 ++- homeassistant/components/unifi/manifest.json | 2 +- homeassistant/components/unifi/switch.py | 104 ++++++++++++++---- requirements_all.txt | 2 +- requirements_test_all.txt | 2 +- tests/components/unifi/test_device_tracker.py | 1 + tests/components/unifi/test_init.py | 2 + tests/components/unifi/test_switch.py | 62 +++++++++-- 10 files changed, 157 insertions(+), 43 deletions(-) diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py index 8568e9f6f75..883e00a5559 100644 --- a/homeassistant/components/unifi/__init__.py +++ b/homeassistant/components/unifi/__init__.py @@ -7,8 +7,8 @@ from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC import homeassistant.helpers.config_validation as cv from .const import ( - CONF_CONTROLLER, CONF_DETECTION_TIME, CONF_SITE_ID, CONF_SSID_FILTER, - CONTROLLER_ID, DOMAIN, UNIFI_CONFIG) + CONF_BLOCK_CLIENT, CONF_CONTROLLER, CONF_DETECTION_TIME, CONF_SITE_ID, + CONF_SSID_FILTER, CONTROLLER_ID, DOMAIN, UNIFI_CONFIG) from .controller import UniFiController CONF_CONTROLLERS = 'controllers' @@ -16,9 +16,11 @@ CONF_CONTROLLERS = 'controllers' CONTROLLER_SCHEMA = vol.Schema({ vol.Required(CONF_HOST): cv.string, vol.Required(CONF_SITE_ID): cv.string, + vol.Optional(CONF_BLOCK_CLIENT, default=[]): vol.All( + cv.ensure_list, [cv.string]), vol.Optional(CONF_DETECTION_TIME): vol.All( cv.time_period, cv.positive_timedelta), - vol.Optional(CONF_SSID_FILTER): vol.All(cv.ensure_list, [cv.string]) + vol.Optional(CONF_SSID_FILTER): vol.All(cv.ensure_list, [cv.string]), }) CONFIG_SCHEMA = vol.Schema({ diff --git a/homeassistant/components/unifi/const.py b/homeassistant/components/unifi/const.py index b638c1664d1..e6076829240 100644 --- a/homeassistant/components/unifi/const.py +++ b/homeassistant/components/unifi/const.py @@ -11,5 +11,6 @@ CONF_SITE_ID = 'site' UNIFI_CONFIG = 'unifi_config' +CONF_BLOCK_CLIENT = 'block_client' CONF_DETECTION_TIME = 'detection_time' CONF_SSID_FILTER = 'ssid_filter' diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index d8d365236d9..ddd0aababac 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -13,7 +13,8 @@ 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) + CONF_BLOCK_CLIENT, CONF_CONTROLLER, CONF_SITE_ID, CONTROLLER_ID, LOGGER, + UNIFI_CONFIG) from .errors import AuthenticationRequired, CannotConnect @@ -52,6 +53,11 @@ class UniFiController: """Return the site user role of this controller.""" return self._site_role + @property + def block_clients(self): + """Return list of clients to block.""" + return self.unifi_config.get(CONF_BLOCK_CLIENT, []) + @property def mac(self): """Return the mac address of this controller.""" @@ -84,6 +90,8 @@ class UniFiController: with async_timeout.timeout(10): await self.api.clients.update() await self.api.devices.update() + if self.block_clients: + await self.api.clients_all.update() except aiounifi.LoginRequired: try: @@ -128,14 +136,14 @@ class UniFiController: except CannotConnect: raise ConfigEntryNotReady - except Exception: # pylint: disable=broad-except + except Exception as err: # pylint: disable=broad-except LOGGER.error( - 'Unknown error connecting with UniFi controller.') + 'Unknown error connecting with UniFi controller: %s', err) 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.site_name == unifi_config[CONF_SITE_ID]: self.unifi_config = unifi_config break diff --git a/homeassistant/components/unifi/manifest.json b/homeassistant/components/unifi/manifest.json index 64119bae2fe..fff731abf4e 100644 --- a/homeassistant/components/unifi/manifest.json +++ b/homeassistant/components/unifi/manifest.json @@ -4,7 +4,7 @@ "config_flow": true, "documentation": "https://www.home-assistant.io/components/unifi", "requirements": [ - "aiounifi==6" + "aiounifi==7" ], "dependencies": [], "codeowners": [ diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 43582e50027..3d3ccc95563 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -51,13 +51,37 @@ def update_items(controller, async_add_entities, switches): new_switches = [] devices = controller.api.devices + # block client + for client_id in controller.block_clients: + + block_client_id = 'block-{}'.format(client_id) + + if block_client_id in switches: + LOGGER.debug("Updating UniFi block switch %s (%s)", + switches[block_client_id].entity_id, + switches[block_client_id].client.mac) + switches[block_client_id].async_schedule_update_ha_state() + continue + + if client_id not in controller.api.clients_all: + continue + + client = controller.api.clients_all[client_id] + switches[block_client_id] = UniFiBlockClientSwitch(client, controller) + new_switches.append(switches[block_client_id]) + LOGGER.debug( + "New UniFi Block switch %s (%s)", client.hostname, client.mac) + + # control poe for client_id in controller.api.clients: - if client_id in switches: - LOGGER.debug("Updating UniFi switch %s (%s)", - switches[client_id].entity_id, - switches[client_id].client.mac) - switches[client_id].async_schedule_update_ha_state() + poe_client_id = 'poe-{}'.format(client_id) + + if poe_client_id in switches: + LOGGER.debug("Updating UniFi POE switch %s (%s)", + switches[poe_client_id].entity_id, + switches[poe_client_id].client.mac) + switches[poe_client_id].async_schedule_update_ha_state() continue client = controller.api.clients[client_id] @@ -80,24 +104,22 @@ def update_items(controller, async_add_entities, switches): if multi_clients_on_port: continue - switches[client_id] = UniFiSwitch(client, controller) - new_switches.append(switches[client_id]) - LOGGER.debug("New UniFi switch %s (%s)", client.hostname, client.mac) + switches[poe_client_id] = UniFiPOEClientSwitch(client, controller) + new_switches.append(switches[poe_client_id]) + LOGGER.debug( + "New UniFi POE switch %s (%s)", client.hostname, client.mac) if new_switches: async_add_entities(new_switches) -class UniFiSwitch(SwitchDevice): - """Representation of a client that uses POE.""" +class UniFiClient: + """Base class for UniFi switches.""" def __init__(self, client, controller): """Set up switch.""" self.client = client self.controller = controller - self.poe_mode = None - if self.port.poe_mode != 'off': - self.poe_mode = self.port.poe_mode async def async_update(self): """Synchronize state with controller.""" @@ -105,8 +127,26 @@ class UniFiSwitch(SwitchDevice): @property def name(self): - """Return the name of the switch.""" - return self.client.hostname + """Return the name of the client.""" + return self.client.name or self.client.hostname + + @property + def device_info(self): + """Return a device description for device registry.""" + return { + 'connections': {(CONNECTION_NETWORK_MAC, self.client.mac)} + } + + +class UniFiPOEClientSwitch(UniFiClient, SwitchDevice): + """Representation of a client that uses POE.""" + + def __init__(self, client, controller): + """Set up POE switch.""" + super().__init__(client, controller) + self.poe_mode = None + if self.port.poe_mode != 'off': + self.poe_mode = self.port.poe_mode @property def unique_id(self): @@ -146,13 +186,6 @@ class UniFiSwitch(SwitchDevice): } return attributes - @property - def device_info(self): - """Return a device description for device registry.""" - return { - 'connections': {(CONNECTION_NETWORK_MAC, self.client.mac)} - } - @property def device(self): """Shortcut to the switch that client is connected to.""" @@ -162,3 +195,30 @@ class UniFiSwitch(SwitchDevice): def port(self): """Shortcut to the switch port that client is connected to.""" return self.device.ports[self.client.sw_port] + + +class UniFiBlockClientSwitch(UniFiClient, SwitchDevice): + """Representation of a blockable client.""" + + @property + def unique_id(self): + """Return a unique identifier for this switch.""" + return 'block-{}'.format(self.client.mac) + + @property + def is_on(self): + """Return true if client is blocked.""" + return self.client.blocked + + @property + def available(self): + """Return if controller is available.""" + return self.controller.available + + async def async_turn_on(self, **kwargs): + """Block client.""" + await self.controller.api.clients.async_block(self.client.mac) + + async def async_turn_off(self, **kwargs): + """Unblock client.""" + await self.controller.api.clients.async_unblock(self.client.mac) diff --git a/requirements_all.txt b/requirements_all.txt index 4e7783e6df9..f2032fde00a 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -169,7 +169,7 @@ aiopvapi==1.6.14 aioswitcher==2019.4.26 # homeassistant.components.unifi -aiounifi==6 +aiounifi==7 # homeassistant.components.wwlln aiowwlln==1.0.0 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index 704c6f7163b..30380c702a5 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -64,7 +64,7 @@ aionotion==1.1.0 aioswitcher==2019.4.26 # homeassistant.components.unifi -aiounifi==6 +aiounifi==7 # homeassistant.components.wwlln aiowwlln==1.0.0 diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index c0b04705e5b..3a209e610ce 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/tests/components/unifi/test_device_tracker.py @@ -122,6 +122,7 @@ async def test_no_clients(hass, mock_controller): """Test the update_clients function when no clients are found.""" mock_controller.mock_client_responses.append({}) mock_controller.mock_device_responses.append({}) + await setup_controller(hass, mock_controller) assert len(mock_controller.mock_requests) == 2 assert len(hass.states.async_all()) == 2 diff --git a/tests/components/unifi/test_init.py b/tests/components/unifi/test_init.py index 27cd55d81f1..2d9ea143e76 100644 --- a/tests/components/unifi/test_init.py +++ b/tests/components/unifi/test_init.py @@ -26,6 +26,7 @@ async def test_setup_with_config(hass): unifi.CONF_CONTROLLERS: { unifi.CONF_HOST: '1.2.3.4', unifi.CONF_SITE_ID: 'My site', + unifi.CONF_BLOCK_CLIENT: ['12:34:56:78:90:AB'], unifi.CONF_DETECTION_TIME: 3, unifi.CONF_SSID_FILTER: ['ssid'] } @@ -36,6 +37,7 @@ async def test_setup_with_config(hass): assert hass.data[unifi.UNIFI_CONFIG] == [{ unifi.CONF_HOST: '1.2.3.4', unifi.CONF_SITE_ID: 'My site', + unifi.CONF_BLOCK_CLIENT: ['12:34:56:78:90:AB'], unifi.CONF_DETECTION_TIME: timedelta(seconds=3), unifi.CONF_SSID_FILTER: ['ssid'] }] diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index 2d64681e161..f467adad9a2 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -7,7 +7,7 @@ import pytest from tests.common import mock_coro import aiounifi -from aiounifi.clients import Clients +from aiounifi.clients import Clients, ClientsAll from aiounifi.devices import Devices from homeassistant import config_entries @@ -170,6 +170,29 @@ DEVICE_1 = { ] } +BLOCKED = { + 'blocked': True, + 'hostname': 'block_client_1', + 'ip': '10.0.0.1', + 'is_guest': False, + 'is_wired': False, + 'mac': '00:00:00:00:01:01', + 'name': 'Block Client 1', + 'noted': True, + 'oui': 'Producer', +} +UNBLOCKED = { + 'blocked': False, + 'hostname': 'block_client_2', + 'ip': '10.0.0.2', + 'is_guest': False, + 'is_wired': True, + 'mac': '00:00:00:00:01:02', + 'name': 'Block Client 2', + 'noted': True, + 'oui': 'Producer', +} + CONTROLLER_DATA = { CONF_HOST: 'mock-host', CONF_USERNAME: 'mock-user', @@ -199,6 +222,7 @@ def mock_controller(hass): controller.mock_client_responses = deque() controller.mock_device_responses = deque() + controller.mock_client_all_responses = deque() async def mock_request(method, path, **kwargs): kwargs['method'] = method @@ -208,10 +232,13 @@ def mock_controller(hass): return controller.mock_client_responses.popleft() if path == 's/{site}/stat/device': return controller.mock_device_responses.popleft() + if path == 's/{site}/rest/user': + return controller.mock_client_all_responses.popleft() return None controller.api.clients = Clients({}, mock_request) controller.api.devices = Devices({}, mock_request) + controller.api.clients_all = ClientsAll({}, mock_request) return controller @@ -277,12 +304,17 @@ async def test_switches(hass, mock_controller): """Test the update_items function with some clients.""" mock_controller.mock_client_responses.append([CLIENT_1, CLIENT_4]) mock_controller.mock_device_responses.append([DEVICE_1]) + mock_controller.mock_client_all_responses.append( + [BLOCKED, UNBLOCKED, CLIENT_1]) + mock_controller.unifi_config = { + unifi.CONF_BLOCK_CLIENT: [BLOCKED['mac'], UNBLOCKED['mac']] + } await setup_controller(hass, mock_controller) - assert len(mock_controller.mock_requests) == 2 - assert len(hass.states.async_all()) == 2 + assert len(mock_controller.mock_requests) == 3 + assert len(hass.states.async_all()) == 4 - switch_1 = hass.states.get('switch.client_1') + switch_1 = hass.states.get('switch.poe_client_1') assert switch_1 is not None assert switch_1.state == 'on' assert switch_1.attributes['power'] == '2.56' @@ -292,9 +324,17 @@ async def test_switches(hass, mock_controller): assert switch_1.attributes['port'] == 1 assert switch_1.attributes['poe_mode'] == 'auto' - switch_4 = hass.states.get('switch.client_4') + switch_4 = hass.states.get('switch.poe_client_4') assert switch_4 is None + blocked = hass.states.get('switch.block_client_1') + assert blocked is not None + assert blocked.state == 'on' + + unblocked = hass.states.get('switch.block_client_2') + assert unblocked is not None + assert unblocked.state == 'off' + async def test_new_client_discovered(hass, mock_controller): """Test if 2nd update has a new client.""" @@ -310,7 +350,7 @@ async def test_new_client_discovered(hass, mock_controller): # Calling a service will trigger the updates to run await hass.services.async_call('switch', 'turn_off', { - 'entity_id': 'switch.client_1' + 'entity_id': 'switch.poe_client_1' }, blocking=True) assert len(mock_controller.mock_requests) == 5 assert len(hass.states.async_all()) == 3 @@ -326,7 +366,7 @@ async def test_new_client_discovered(hass, mock_controller): } await hass.services.async_call('switch', 'turn_on', { - 'entity_id': 'switch.client_1' + 'entity_id': 'switch.poe_client_1' }, blocking=True) assert len(mock_controller.mock_requests) == 7 assert mock_controller.mock_requests[5] == { @@ -340,7 +380,7 @@ async def test_new_client_discovered(hass, mock_controller): 'path': 's/{site}/rest/device/mock-id' } - switch_2 = hass.states.get('switch.client_2') + switch_2 = hass.states.get('switch.poe_client_2') assert switch_2 is not None assert switch_2.state == 'on' @@ -384,7 +424,7 @@ async def test_failed_update_unreachable_controller(hass, mock_controller): # Calling a service will trigger the updates to run await hass.services.async_call('switch', 'turn_off', { - 'entity_id': 'switch.client_1' + 'entity_id': 'switch.poe_client_1' }, blocking=True) assert len(mock_controller.mock_requests) == 3 @@ -406,7 +446,7 @@ async def test_ignore_multiple_poe_clients_on_same_port(hass, mock_controller): # 1 All Lights group, 2 lights assert len(hass.states.async_all()) == 0 - switch_1 = hass.states.get('switch.client_1') - switch_2 = hass.states.get('switch.client_2') + switch_1 = hass.states.get('switch.poe_client_1') + switch_2 = hass.states.get('switch.poe_client_2') assert switch_1 is None assert switch_2 is None