UniFi block clients (#25478)

* Allow blocking clients on UniFi networks

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)