Home-Assistant
/
core
mirror of https://github.com/home-assistant/core.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

UniFi block clients (#25478) 

* Allow blocking clients on UniFi networks
pull/25485/head
Robert Svensson 2019-07-25 16:56:56 +02:00 committed by GitHub
parent 59c62a261b
commit b6934f0cd0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
10 changed files with 157 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
homeassistant/components/unifi/__init__.py
View File

@ -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({

1
homeassistant/components/unifi/const.py
View File

@ -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'

16
homeassistant/components/unifi/controller.py
View File

@ -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

2
homeassistant/components/unifi/manifest.json
View File

@ -4,7 +4,7 @@
"config_flow": true,
"documentation": "https://www.home-assistant.io/components/unifi",
"requirements": [
"aiounifi==6"
"aiounifi==7"
],
"dependencies": [],
"codeowners": [

104
homeassistant/components/unifi/switch.py
View File

@ -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)

2
requirements_all.txt
View File

@ -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

2
requirements_test_all.txt
View File

@ -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

1
tests/components/unifi/test_device_tracker.py
View File

@ -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

2
tests/components/unifi/test_init.py
View File

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

62
tests/components/unifi/test_switch.py
View File

@ -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