Device tags added (#9)

vmakeev · web-flow · commit 715154976301 · 2022-11-14T00:21:21.000+04:00
diff --git a/README.md b/README.md
@@ -2,8 +2,9 @@
 
 Home Assistant custom component for control [Huawei WiFi Mesh 3](https://consumer.huawei.com/en/routers/wifi-mesh3/) routers over LAN.
 
-**0.7.3**
+**0.7.4**
 
+- tagging connected devices
 - sensors for the number of connected devices (total and for each individual router)
 - enable/disable NFC on each router separately
 - enable/disable TWT (reduce power consumption of Wi-Fi 6 devices in sleep mode)
@@ -38,6 +39,7 @@ Each tracked device exposes the following attributes:
 | `rssi`           | Signal strength for wireless connections     | Yes                 |
 | `is_guest`       | Is the device connected to the guest network | Yes                 |
 | `is_hilink`      | Is the device connected via HiLink (usually other routers) | Yes   |
+| `tags`           | List of [tags](#device-tags) that marked the device          | No                  |
 | `friendly_name`  | Device name provided by the router           | No                  |
 
 Tracked device names, including routers, can be changed in [your mesh control interface](http://192.168.3.1/html/index.html#/devicecontrol), after which the component will update them in Home Assistant
@@ -71,11 +73,84 @@ _Note: when additional routers are disconnected from the network, their personal
 
 Each sensor exposes the following attributes:
 
-|     Attribute      |                  Description                     |
-|--------------------|--------------------------------------------------|
-| `guest_clients`    | Number of devices connected to the guest network |
-| `hilink_clients`   | Number of devices connected via HiLink           |
-| `wireless_clients` | Number of devices connected wirelessly           |
-| `lan_clients`      | Number of devices connected by cable             |
-| `wifi_2_4_clients` | Number of devices connected to Wi-Fi 2.4 GHz     |
-| `wifi_5_clients`   | Number of devices connected to Wi-Fi 5 GHz       |
+|         Attribute            |                         Description                          |
+|------------------------------|--------------------------------------------------------------|
+| `guest_clients`              | Number of devices connected to the guest network             |
+| `hilink_clients`             | Number of devices connected via HiLink                       |
+| `wireless_clients`           | Number of devices connected wirelessly                       |
+| `lan_clients`                | Number of devices connected by cable                         |
+| `wifi_2_4_clients`           | Number of devices connected to Wi-Fi 2.4 GHz                 |
+| `wifi_5_clients`             | Number of devices connected to Wi-Fi 5 GHz                   |
+| `tagged_<tag_name>_clients`  | Number of connected devices with a specific [tag](#device-tags) `<tag_name>` |
+
+## Customization
+
+### Device tags
+
+The component allows you to attach one or more tags to each client device in order to be able to use in automation the number of devices marked with a tag, connected to a specific router, or to the entire mesh network.
+
+The component will attempt to load the device tag-to-MAC mapping from the file located at `<home assistant config folder>/.storage/huawei_mesh_<long_config_id>_tags`. If the file does not exist, then the component will create it with a usage example:
+
+```
+{
+  "version": 1,
+  "minor_version": 1,
+  "key": "huawei_mesh_<long_config_id>_tags",
+  "data": {
+    "homeowners": [
+      "place_mac_addresses_here"
+    ],
+    "visitors": [
+      "place_mac_addresses_here"
+    ]
+  }
+}
+```
+
+_Note: unfortunately, editing the list of tags and devices associated with them is currently available only through editing this file._
+
+Each tag can have multiple devices associated with it. Each device can be associated with multiple tags.
+
+Example:
+```
+{
+  "version": 1,
+  "minor_version": 1,
+  "key": "huawei_mesh_<long_config_id>_tags",
+  "data": {
+    "my_awesome_tag": [
+      "00:11:22:33:44:55",
+      "A0:B1:C2:D3:E4:F5",
+      "F5:E4:D3:C2:B1:A0"
+    ],
+    "another_tag": [
+      "00:11:22:33:44:55",
+      "A9:B8:C7:D6:E5:F4"
+    ],
+    "third_tag": [
+      "99:88:77:66:55:44"
+    ]
+  }
+}
+```
+
+**Usage example:**
+
+|  Tag name    |          Tagged Devices           |
+|--------------|-----------------------------------|
+|  homeowners  | Michael's phone, Michael's laptop |
+|  visitors    | Victoria's phone, Eugene's phone  |
+
+
+- Michael's phone is connected to the "Garage" router
+- Michael's laptop is connected to the "Living room" router
+- Victoria's phone is connected to the "Living room" router
+- Eugene's phone is connected to the "primary" router
+
+In this scenario, the sensors for the number of connected devices will provide the following attributes:
+
+|                 Sensor                        | Attributes and values |
+|-----------------------------------------------|-----------------------|
+| `sensor.huawei_mesh_3_clients_garage`         | `guest_clients`: 0 <br/> `hilink_clients`: 0<br/>`wireless_clients`: 1<br />`lan_clients`: 0<br />`wifi_2_4_clients`: 0<br />`wifi_5_clients`: 1<br />`tagged_homeowners_clients`: 1 _// Michael's phone_<br />`tagged_visitors_clients`: 0 |
+| `sensor.huawei_mesh_3_clients_living_room`    | `guest_clients`: 0 <br/> `hilink_clients`: 0<br/>`wireless_clients`: 2<br />`lan_clients`: 0<br />`wifi_2_4_clients`: 0<br />`wifi_5_clients`: 2<br />`tagged_homeowners_clients`: 1 _// Michael's laptop_<br />`tagged_visitors_clients`: 1 _// Victoria's phone_ |
+| `sensor.huawei_mesh_3_clients_primary_router` | `guest_clients`: 0 <br/> `hilink_clients`: 2<br/>`wireless_clients`: 3<br />`lan_clients`: 0<br />`wifi_2_4_clients`: 0<br />`wifi_5_clients`: 3<br />`tagged_homeowners_clients`: 0<br />`tagged_visitors_clients`: 1 _// Eugene's phone_ |
diff --git a/custom_components/huawei_mesh_router/__init__.py b/custom_components/huawei_mesh_router/__init__.py
@@ -4,8 +4,10 @@
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import config_validation
+from homeassistant.helpers.storage import Store
 
 from .const import (
+    STORAGE_VERSION,
     DOMAIN,
     PLATFORMS,
     ATTR_MANUFACTURER
@@ -31,7 +33,8 @@ async def async_setup(hass, _config):
 # ---------------------------
 async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
     """Set up Huawei Router as config entry."""
-    coordinator = HuaweiControllerDataUpdateCoordinator(hass, config_entry)
+    store: Store = Store(hass, STORAGE_VERSION, f"huawei_mesh_{config_entry.entry_id}_tags")
+    coordinator = HuaweiControllerDataUpdateCoordinator(hass, config_entry, store)
     await coordinator.async_config_entry_first_refresh()
 
     config_entry.async_on_unload(config_entry.add_update_listener(update_listener))
diff --git a/custom_components/huawei_mesh_router/connected_device.py b/custom_components/huawei_mesh_router/connected_device.py
@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Any, Dict, Iterable, Tuple
 from homeassistant.backports.enum import StrEnum
 
 from .const import VENDOR_CLASS_ID_ROUTER
@@ -23,21 +23,25 @@ def __init__(self,
                  host_name: str,
                  mac: str,
                  is_active: bool,
+                 tags: list[str],
                  **kwargs: Dict) -> None:
         self._name: str = name
         self._host_name: str = host_name
         self._mac: str = mac
         self._is_active: bool = is_active
+        self._tags: list[str] = tags
         self._data: Dict = kwargs or {}
 
     def update_device_data(self,
                            name: str,
                            host_name: str,
                            is_active: bool,
+                           tags: list[str],
                            **kwargs: Dict):
         self._name: str = name
         self._host_name: str = host_name
         self._is_active: bool = is_active
+        self._tags: list[str] = tags
         self._data: Dict = kwargs or {}
 
     def __str__(self) -> str:
@@ -97,6 +101,13 @@ def is_router(self) -> bool:
         return self.is_hilink and self._data.get("vendor_class_id") == VENDOR_CLASS_ID_ROUTER
 
     @property
-    def all_attrs(self) -> Dict:
+    def tags(self) -> list[str]:
+        """Return device tags list."""
+        return self._tags
+
+    @property
+    def all_attrs(self) -> Iterable[Tuple[str, Any]]:
         """Return dictionary with additional attributes."""
-        return self._data
+        for key, value in self._data.items():
+            yield key, value
+        yield "tags", self._tags
diff --git a/custom_components/huawei_mesh_router/const.py b/custom_components/huawei_mesh_router/const.py
@@ -2,6 +2,8 @@
 
 DOMAIN = "huawei_mesh_router"
 
+STORAGE_VERSION = 1
+
 DEFAULT_HOST = "192.168.3.1"
 DEFAULT_USER = "admin"
 DEFAULT_PORT = 80
diff --git a/custom_components/huawei_mesh_router/device_tracker.py b/custom_components/huawei_mesh_router/device_tracker.py
@@ -107,7 +107,7 @@ def unique_id(self) -> str:
     @property
     def extra_state_attributes(self) -> dict[str, Any]:
         """Return the device state attributes."""
-        return {k: v for k, v in self.device.all_attrs.items() if k not in FILTER_ATTRS}
+        return {k: v for k, v in self.device.all_attrs if k not in FILTER_ATTRS}
 
     @property
     def entity_registry_enabled_default(self) -> bool:
diff --git a/custom_components/huawei_mesh_router/manifest.json b/custom_components/huawei_mesh_router/manifest.json
@@ -1,9 +1,11 @@
 {
   "domain": "huawei_mesh_router",
   "name": "Huawei Mesh Router",
+  "integration_type": "device",
   "documentation": "https://github.com/vmakeev/huawei_mesh_router",
+  "issue_tracker": "https://github.com/vmakeev/huawei_mesh_router/issues",
   "codeowners": ["@vmakeev"],
   "iot_class": "local_polling",
-  "version": "0.7.3",
+  "version": "0.7.4",
   "config_flow": true
 }
diff --git a/custom_components/huawei_mesh_router/sensor.py b/custom_components/huawei_mesh_router/sensor.py
@@ -6,7 +6,6 @@
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import callback, HomeAssistant
-from homeassistant.helpers import entity_registry
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.entity_registry import EntityRegistry
@@ -160,6 +159,10 @@ def _handle_coordinator_update(self) -> None:
         wifi_2_4_clients: int = 0
         wifi_5_clients: int = 0
 
+        tagged_devices: dict[str, int] = {}
+        for tag in self.coordinator.tags_map.get_all_tags():
+            tagged_devices[tag] = 0
+
         for device in self.coordinator.connected_devices.values():
 
             if not self._devices_predicate(device):
@@ -181,6 +184,12 @@ def _handle_coordinator_update(self) -> None:
                 wireless_clients += 1
                 wifi_5_clients += 1
 
+            for tag in device.tags:
+                if tag in tagged_devices:
+                    tagged_devices[tag] += 1
+                else:
+                    tagged_devices[tag] = 1
+
         self._actual_value = total_clients
 
         self._attrs["guest_clients"] = guest_clients
@@ -190,6 +199,9 @@ def _handle_coordinator_update(self) -> None:
         self._attrs["wifi_2_4_clients"] = wifi_2_4_clients
         self._attrs["wifi_5_clients"] = wifi_5_clients
 
+        for tag, count in tagged_devices.items():
+            self._attrs[f"tagged_{tag}_clients"] = count
+
         super()._handle_coordinator_update()
 
     @property
diff --git a/custom_components/huawei_mesh_router/update_coordinator.py b/custom_components/huawei_mesh_router/update_coordinator.py
@@ -1,10 +1,9 @@
 """Huawei Controller for Huawei Router."""
 from __future__ import annotations
 
-import asyncio
 import logging
 from datetime import timedelta
-from typing import Any, cast, Callable, Iterator, List
+from typing import Any, cast, Callable, Iterable, List
 
 from aiohttp import ClientResponse
 
@@ -13,6 +12,7 @@
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.helpers import entity_registry
 from homeassistant.helpers.entity_registry import EntityRegistry
+from homeassistant.helpers.storage import Store
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 from .connected_device import ConnectedDevice
 from .router_info import RouterInfo
@@ -107,13 +107,62 @@ def watch_for_changes(
                 on_router_removed(er, device_id, unavailable_router)
 
 
+# ---------------------------
+#   TagsMap
+# ---------------------------
+class TagsMap:
+
+    def __init__(self, tags_map_storage: Store):
+        self._storage: Store = tags_map_storage
+        self._mac_to_tags: dict[str, list[str]] = {}
+        self._tag_to_macs: dict[str, list[str]] = {}
+        self._is_loaded: bool = False
+
+    @property
+    def is_loaded(self) -> bool:
+        return self._is_loaded
+
+    async def load(self):
+        _LOGGER.debug("Stored tags loading started")
+
+        self._mac_to_tags.clear()
+        self._tag_to_macs.clear()
+
+        self._tag_to_macs = await self._storage.async_load()
+        if not self._tag_to_macs:
+            _LOGGER.debug("No stored tags found, creating sample")
+            default_tags = {"homeowners": ["place_mac_addresses_here"], "visitors": ["place_mac_addresses_here"]}
+            await self._storage.async_save(default_tags)
+            self._tag_to_macs = default_tags
+
+        for tag, devices_macs in self._tag_to_macs.items():
+            for device_mac in devices_macs:
+                if device_mac not in self._mac_to_tags:
+                    self._mac_to_tags[device_mac] = []
+                self._mac_to_tags[device_mac].append(tag)
+
+        self._is_loaded = True
+
+        _LOGGER.debug("Stored tags loading finished")
+
+    def get_tags(self, mac_address: str) -> list[str]:
+        return self._mac_to_tags.get(mac_address, [])
+
+    def get_all_tags(self) -> Iterable[str]:
+        return self._tag_to_macs.keys()
+
+    def get_devices(self, tag: str) -> list[str]:
+        return self._tag_to_macs.get(tag, [])
+
+
 # ---------------------------
 #   HuaweiControllerDataUpdateCoordinator
 # ---------------------------
 class HuaweiControllerDataUpdateCoordinator(DataUpdateCoordinator):
 
-    def __init__(self, hass: HomeAssistant, config_entry: ConfigEntry) -> None:
+    def __init__(self, hass: HomeAssistant, config_entry: ConfigEntry, tags_map_storage: Store) -> None:
         """Initialize HuaweiController."""
+        self._tags_map: TagsMap = TagsMap(tags_map_storage)
         self._router_info: RouterInfo | None = None
         self._switch_states: dict[str, bool] = {}
         self._connected_devices: dict[str, ConnectedDevice] = {}
@@ -188,6 +237,10 @@ def device_info(self) -> DeviceInfo:
             sw_version=self.router_info.software_version,
         )
 
+    @property
+    def tags_map(self) -> TagsMap:
+        return self._tags_map
+
     def _safe_disconnect(self, api: HuaweiApi) -> None:
         """Disconnect from API."""
         try:
@@ -289,7 +342,7 @@ async def _update_connected_devices(self) -> None:
 
         """recursively search all HiLink routers with connected devices"""
 
-        def get_mesh_routers(devices: List[dict[str, Any]]) -> Iterator[dict[str, Any]]:
+        def get_mesh_routers(devices: List[dict[str, Any]]) -> Iterable[dict[str, Any]]:
             for candidate in devices:
                 if candidate.get('HiLinkType') == "Device":
                     yield candidate
@@ -316,16 +369,20 @@ def get_mesh_routers(devices: List[dict[str, Any]]) -> Iterator[dict[str, Any]]:
                         "id": router.get('MACAddress')
                     }
 
+        if not self._tags_map.is_loaded:
+            await self._tags_map.load()
+
         for device_data in devices_data:
             mac: str = device_data['MACAddress']
             host_name: str = device_data.get('HostName', f'device_{mac}')
             name: str = device_data.get('ActualName', host_name)
             is_active: bool = device_data.get('Active', False)
+            tags = self._tags_map.get_tags(mac)
 
             if mac in self._connected_devices:
                 device = self._connected_devices[mac]
             else:
-                device = ConnectedDevice(name, host_name, mac, is_active)
+                device = ConnectedDevice(name, host_name, mac, is_active, tags)
                 self._connected_devices[device.mac] = device
 
             if is_active:
@@ -335,7 +392,7 @@ def get_mesh_routers(devices: List[dict[str, Any]]) -> Iterator[dict[str, Any]]:
                                                            "name": self.name or 'Primary router',
                                                            "id": CONNECTED_VIA_ID_PRIMARY
                                                        })
-                device.update_device_data(name, host_name, True,
+                device.update_device_data(name, host_name, True, tags,
                                           connected_via=connected_via.get("name"),
                                           ip_address=device_data.get('IPAddress'),
                                           interface_type=device_data.get('InterfaceType'),
@@ -346,7 +403,7 @@ def get_mesh_routers(devices: List[dict[str, Any]]) -> Iterator[dict[str, Any]]:
                                           connected_via_id=connected_via.get("id")
                                           )
             else:
-                device.update_device_data(name, host_name, False)
+                device.update_device_data(name, host_name, False, tags)
 
         _LOGGER.debug('Connected devices updated')
 

Original file line number	Diff line number	Diff line change
`@@ -1,9 +1,11 @@`
`1`	`1`	`{`
`2`	`2`	`"domain": "huawei_mesh_router",`
`3`	`3`	`"name": "Huawei Mesh Router",`
	`4`	`+ "integration_type": "device",`
`4`	`5`	`"documentation": "https://github.com/vmakeev/huawei_mesh_router",`
	`6`	`+ "issue_tracker": "https://github.com/vmakeev/huawei_mesh_router/issues",`
`5`	`7`	`"codeowners": ["@vmakeev"],`
`6`	`8`	`"iot_class": "local_polling",`
`7`		`- "version": "0.7.3",`
	`9`	`+ "version": "0.7.4",`
`8`	`10`	`"config_flow": true`
`9`	`11`	`}`