UniFi - Add block network access control to config option (home-assis…

…tant#32004)

* Add block network access control to config option

* Clean up

homeassistant/components/unifi/.translations/en.json

@@ -6,7 +6,8 @@
         },
         "error": {
             "faulty_credentials": "Bad user credentials",
-            "service_unavailable": "No service available"
+            "service_unavailable": "No service available",
+            "unknown_client_mac": "No client available on that MAC address"
         },
         "step": {
             "user": {
@@ -34,15 +35,26 @@
                     "track_wired_clients": "Include wired network clients"
                 },
                 "description": "Configure device tracking",
-                "title": "UniFi options"
+                "title": "UniFi options 1/3"
+            },
+            "client_control": {
+                "data": {
+                    "block_client": "Network access controlled clients",
+                    "new_client": "Add new client (MAC) for network access control"
+                },
+                "description": "Configure client controls\n\nCreate switches for serial numbers you want to control network access for.",
+                "title": "UniFi options 2/3"
             },
             "statistics_sensors": {
                 "data": {
                     "allow_bandwidth_sensors": "Bandwidth usage sensors for network clients"
                 },
                 "description": "Configure statistics sensors",
-                "title": "UniFi options"
+                "title": "UniFi options 3/3"
             }
+        },
+        "error": {
+            "unknown_client_mac": "No client available in UniFi on that MAC address"
         }
     }
 }

homeassistant/components/unifi/config_flow.py

@@ -16,6 +16,7 @@
 
 from .const import (
     CONF_ALLOW_BANDWIDTH_SENSORS,
+    CONF_BLOCK_CLIENT,
     CONF_CONTROLLER,
     CONF_DETECTION_TIME,
     CONF_SITE_ID,
@@ -30,6 +31,7 @@
 from .controller import get_controller
 from .errors import AlreadyConfigured, AuthenticationRequired, CannotConnect
 
+CONF_NEW_CLIENT = "new_client"
 DEFAULT_PORT = 8443
 DEFAULT_SITE_ID = "default"
 DEFAULT_VERIFY_SSL = False
@@ -171,61 +173,117 @@ def __init__(self, config_entry):
         """Initialize UniFi options flow."""
         self.config_entry = config_entry
         self.options = dict(config_entry.options)
+        self.controller = None
 
     async def async_step_init(self, user_input=None):
         """Manage the UniFi options."""
+        self.controller = get_controller_from_config_entry(self.hass, self.config_entry)
+        self.options[CONF_BLOCK_CLIENT] = self.controller.option_block_clients
         return await self.async_step_device_tracker()
 
     async def async_step_device_tracker(self, user_input=None):
         """Manage the device tracker options."""
         if user_input is not None:
             self.options.update(user_input)
-            return await self.async_step_statistics_sensors()
+            return await self.async_step_client_control()
 
-        controller = get_controller_from_config_entry(self.hass, self.config_entry)
-
-        ssid_filter = {wlan: wlan for wlan in controller.api.wlans}
+        ssid_filter = {wlan: wlan for wlan in self.controller.api.wlans}
 
         return self.async_show_form(
             step_id="device_tracker",
             data_schema=vol.Schema(
                 {
                     vol.Optional(
-                        CONF_TRACK_CLIENTS, default=controller.option_track_clients,
+                        CONF_TRACK_CLIENTS,
+                        default=self.controller.option_track_clients,
                     ): bool,
                     vol.Optional(
                         CONF_TRACK_WIRED_CLIENTS,
-                        default=controller.option_track_wired_clients,
+                        default=self.controller.option_track_wired_clients,
                     ): bool,
                     vol.Optional(
-                        CONF_TRACK_DEVICES, default=controller.option_track_devices,
+                        CONF_TRACK_DEVICES,
+                        default=self.controller.option_track_devices,
                     ): bool,
                     vol.Optional(
-                        CONF_SSID_FILTER, default=controller.option_ssid_filter
+                        CONF_SSID_FILTER, default=self.controller.option_ssid_filter
                     ): cv.multi_select(ssid_filter),
                     vol.Optional(
                         CONF_DETECTION_TIME,
-                        default=int(controller.option_detection_time.total_seconds()),
+                        default=int(
+                            self.controller.option_detection_time.total_seconds()
+                        ),
                     ): int,
                 }
             ),
         )
 
+    async def async_step_client_control(self, user_input=None):
+        """Manage configuration of network access controlled clients."""
+        errors = {}
+
+        if user_input is not None:
+            new_client = user_input.pop(CONF_NEW_CLIENT, None)
+            self.options.update(user_input)
+
+            if new_client:
+                if (
+                    new_client in self.controller.api.clients
+                    or new_client in self.controller.api.clients_all
+                ):
+                    self.options[CONF_BLOCK_CLIENT].append(new_client)
+
+                else:
+                    errors["base"] = "unknown_client_mac"
+
+            else:
+                return await self.async_step_statistics_sensors()
+
+        clients_to_block = {}
+
+        for mac in self.options[CONF_BLOCK_CLIENT]:
+
+            name = None
+
+            for clients in [
+                self.controller.api.clients,
+                self.controller.api.clients_all,
+            ]:
+                if mac in clients:
+                    name = f"{clients[mac].name or clients[mac].hostname} ({mac})"
+                    break
+
+            if not name:
+                name = mac
+
+            clients_to_block[mac] = name
+
+        return self.async_show_form(
+            step_id="client_control",
+            data_schema=vol.Schema(
+                {
+                    vol.Optional(
+                        CONF_BLOCK_CLIENT, default=self.options[CONF_BLOCK_CLIENT]
+                    ): cv.multi_select(clients_to_block),
+                    vol.Optional(CONF_NEW_CLIENT): str,
+                }
+            ),
+            errors=errors,
+        )
+
     async def async_step_statistics_sensors(self, user_input=None):
         """Manage the statistics sensors options."""
         if user_input is not None:
             self.options.update(user_input)
             return await self._update_options()
 
-        controller = get_controller_from_config_entry(self.hass, self.config_entry)
-
         return self.async_show_form(
             step_id="statistics_sensors",
             data_schema=vol.Schema(
                 {
                     vol.Optional(
                         CONF_ALLOW_BANDWIDTH_SENSORS,
-                        default=controller.option_allow_bandwidth_sensors,
+                        default=self.controller.option_allow_bandwidth_sensors,
                     ): bool
                 }
             ),

homeassistant/components/unifi/const.py

@@ -25,11 +25,9 @@
 CONF_DONT_TRACK_WIRED_CLIENTS = "dont_track_wired_clients"
 
 DEFAULT_ALLOW_BANDWIDTH_SENSORS = False
-DEFAULT_BLOCK_CLIENTS = []
 DEFAULT_TRACK_CLIENTS = True
 DEFAULT_TRACK_DEVICES = True
 DEFAULT_TRACK_WIRED_CLIENTS = True
 DEFAULT_DETECTION_TIME = 300
-DEFAULT_SSID_FILTER = []
 
 ATTR_MANUFACTURER = "Ubiquiti Networks"

homeassistant/components/unifi/controller.py

@@ -31,9 +31,7 @@
     CONF_TRACK_WIRED_CLIENTS,
     CONTROLLER_ID,
     DEFAULT_ALLOW_BANDWIDTH_SENSORS,
-    DEFAULT_BLOCK_CLIENTS,
     DEFAULT_DETECTION_TIME,
-    DEFAULT_SSID_FILTER,
     DEFAULT_TRACK_CLIENTS,
     DEFAULT_TRACK_DEVICES,
     DEFAULT_TRACK_WIRED_CLIENTS,
@@ -99,7 +97,7 @@ def option_allow_bandwidth_sensors(self):
     @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, DEFAULT_BLOCK_CLIENTS)
+        return self.config_entry.options.get(CONF_BLOCK_CLIENT, [])
 
     @property
     def option_track_clients(self):
@@ -130,7 +128,7 @@ def option_detection_time(self):
     @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, DEFAULT_SSID_FILTER)
+        return self.config_entry.options.get(CONF_SSID_FILTER, [])
 
     @property
     def mac(self):

homeassistant/components/unifi/strings.json

@@ -16,7 +16,8 @@
         },
         "error": {
             "faulty_credentials": "Bad user credentials",
-            "service_unavailable": "No service available"
+            "service_unavailable": "No service available",
+            "unknown_client_mac": "No client available on that MAC address"
         },
         "abort": {
             "already_configured": "Controller site is already configured",
@@ -37,15 +38,26 @@
                     "track_wired_clients": "Include wired network clients"
                 },
                 "description": "Configure device tracking",
-                "title": "UniFi options"
+                "title": "UniFi options 1/3"
+            },
+            "client_control": {
+                "data": {
+                    "block_client": "Network access controlled clients",
+                    "new_client": "Add new client for network access control"
+                },
+                "description": "Configure client controls\n\nCreate switches for serial numbers you want to control network access for.",
+                "title": "UniFi options 2/3"
             },
             "statistics_sensors": {
                 "data": {
                     "allow_bandwidth_sensors": "Bandwidth usage sensors for network clients"
                 },
                 "description": "Configure statistics sensors",
-                "title": "UniFi options"
+                "title": "UniFi options 3/3"
             }
         }
+    },
+    "error": {
+        "unknown_client_mac": "No client available in UniFi on that MAC address"
     }
 }

homeassistant/components/unifi/switch.py

@@ -4,7 +4,6 @@
 from homeassistant.components.switch import SwitchDevice
 from homeassistant.components.unifi.config_flow import get_controller_from_config_entry
 from homeassistant.core import callback
-from homeassistant.helpers import entity_registry
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.restore_state import RestoreEntity
 
@@ -30,10 +29,12 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     switches = {}
     switches_off = []
 
-    registry = await entity_registry.async_get_registry(hass)
+    option_block_clients = controller.option_block_clients
+
+    entity_registry = await hass.helpers.entity_registry.async_get_registry()
 
     # Restore clients that is not a part of active clients list.
-    for entity in registry.entities.values():
+    for entity in entity_registry.entities.values():
 
         if (
             entity.config_entry_id == config_entry.entry_id
@@ -61,6 +62,43 @@ def update_controller():
         async_dispatcher_connect(hass, controller.signal_update, update_controller)
     )
 
+    @callback
+    def options_updated():
+        """Manage entities affected by config entry options."""
+        nonlocal option_block_clients
+
+        update = set()
+        remove = set()
+
+        if option_block_clients != controller.option_block_clients:
+            option_block_clients = controller.option_block_clients
+
+            for block_client_id, entity in switches.items():
+                if not isinstance(entity, UniFiBlockClientSwitch):
+                    continue
+
+                if entity.client.mac in option_block_clients:
+                    update.add(block_client_id)
+                else:
+                    remove.add(block_client_id)
+
+            for block_client_id in remove:
+                entity = switches.pop(block_client_id)
+
+                if entity_registry.async_is_registered(entity.entity_id):
+                    entity_registry.async_remove(entity.entity_id)
+
+                hass.async_create_task(entity.async_remove())
+
+            if len(update) != len(option_block_clients):
+                update_controller()
+
+    controller.listeners.append(
+        async_dispatcher_connect(
+            hass, controller.signal_options_update, options_updated
+        )
+    )
+
     update_controller()
     switches_off.clear()
 
@@ -74,15 +112,21 @@ def add_entities(controller, async_add_entities, switches, switches_off):
     # block client
     for client_id in controller.option_block_clients:
 
+        client = None
         block_client_id = f"block-{client_id}"
 
         if block_client_id in switches:
             continue
 
-        if client_id not in controller.api.clients_all:
+        if client_id in controller.api.clients:
+            client = controller.api.clients[client_id]
+
+        elif client_id in controller.api.clients_all:
+            client = controller.api.clients_all[client_id]
+
+        if not client:
             continue
 
-        client = controller.api.clients_all[client_id]
         switches[block_client_id] = UniFiBlockClientSwitch(client, controller)
         new_switches.append(switches[block_client_id])