UniFi block clients (#25478)

* Allow blocking clients on UniFi networks

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,
+    CONF_BLOCK_CLIENT, CONF_CONTROLLER, CONF_DETECTION_TIME, CONF_SITE_ID,
-    CONTROLLER_ID, DOMAIN, UNIFI_CONFIG)
+    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({

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'

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
 

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": [

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:
+        poe_client_id = 'poe-{}'.format(client_id)
-            LOGGER.debug("Updating UniFi switch %s (%s)",
+
-                         switches[client_id].entity_id,
+        if poe_client_id in switches:
-                         switches[client_id].client.mac)
+            LOGGER.debug("Updating UniFi POE switch %s (%s)",
-            switches[client_id].async_schedule_update_ha_state()
+                         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)
+        switches[poe_client_id] = UniFiPOEClientSwitch(client, controller)
-        new_switches.append(switches[client_id])
+        new_switches.append(switches[poe_client_id])
-        LOGGER.debug("New UniFi switch %s (%s)", client.hostname, client.mac)
+        LOGGER.debug(
+            "New UniFi POE switch %s (%s)", client.hostname, client.mac)
 
     if new_switches:
         async_add_entities(new_switches)
 
 
-class UniFiSwitch(SwitchDevice):
+class UniFiClient:
-    """Representation of a client that uses POE."""
+    """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 the name of the client."""
-        return self.client.hostname
+        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)

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

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

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

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']
     }]

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(mock_controller.mock_requests) == 3
-    assert len(hass.states.async_all()) == 2
+    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_1 = hass.states.get('switch.poe_client_1')
-    switch_2 = hass.states.get('switch.client_2')
+    switch_2 = hass.states.get('switch.poe_client_2')
     assert switch_1 is None
     assert switch_2 is None