From c7d5639e84c55cdfddf4decba428c3004f4924e6 Mon Sep 17 00:00:00 2001 From: Snuffy2 Date: Sun, 29 Sep 2024 17:25:33 -0400 Subject: [PATCH] Optimize Device Trackers --- custom_components/opnsense/__init__.py | 3 +- custom_components/opnsense/device_tracker.py | 477 ++++++++++--------- custom_components/opnsense/manifest.json | 2 +- requirements_dev.txt | 2 - 4 files changed, 244 insertions(+), 240 deletions(-) diff --git a/custom_components/opnsense/__init__.py b/custom_components/opnsense/__init__.py index 2b9e079..bf12bf5 100644 --- a/custom_components/opnsense/__init__.py +++ b/custom_components/opnsense/__init__.py @@ -20,7 +20,6 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.entity_registry import async_get from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue -from homeassistant.helpers.restore_state import RestoreEntity from homeassistant.helpers.typing import ConfigType from homeassistant.helpers.update_coordinator import CoordinatorEntity from homeassistant.util import slugify @@ -259,7 +258,7 @@ async def async_remove_entity(self, entity): registry.async_remove(entity.entity_id) -class OPNsenseEntity(CoordinatorEntity[OPNsenseDataUpdateCoordinator], RestoreEntity): +class OPNsenseEntity(CoordinatorEntity[OPNsenseDataUpdateCoordinator]): """Base entity for OPNsense""" def __init__( diff --git a/custom_components/opnsense/device_tracker.py b/custom_components/opnsense/device_tracker.py index 633c20d..02822db 100644 --- a/custom_components/opnsense/device_tracker.py +++ b/custom_components/opnsense/device_tracker.py @@ -1,7 +1,6 @@ -"""Support for tracking for pfSense devices.""" - -from __future__ import annotations +"""Support for tracking for OPNsense devices.""" +from datetime import datetime, timedelta import logging import time from typing import Any, Mapping @@ -16,14 +15,15 @@ async_get as async_get_dev_reg, ) from homeassistant.helpers.entity import DeviceInfo -from homeassistant.util import slugify -from mac_vendor_lookup import AsyncMacLookup +from homeassistant.helpers.restore_state import RestoreEntity -from . import CoordinatorEntityManager, OPNsenseEntity +from . import OPNsenseEntity from .const import ( CONF_DEVICE_TRACKER_CONSIDER_HOME, + CONF_DEVICE_TRACKER_ENABLED, CONF_DEVICES, DEFAULT_DEVICE_TRACKER_CONSIDER_HOME, + DEFAULT_DEVICE_TRACKER_ENABLED, DEVICE_TRACKER_COORDINATOR, DOMAIN, SHOULD_RELOAD, @@ -35,257 +35,196 @@ _LOGGER: logging.Logger = logging.getLogger(__name__) -def lookup_mac(mac_vendor_lookup: AsyncMacLookup, mac: str) -> str: - mac = mac_vendor_lookup.sanitise(mac) - if type(mac) == str: - mac = mac.encode("utf8") - return mac_vendor_lookup.prefixes[mac[:6]].decode("utf8") - - -def get_device_tracker_unique_id(mac: str, device_id: str): - """Generate device_tracker unique ID.""" - return slugify(f"{device_id}_mac_{mac}") - - async def async_setup_entry( hass: HomeAssistant, config_entry: ConfigEntry, async_add_entities: entity_platform.AddEntitiesCallback, ) -> None: - """Set up device tracker for pfSense component.""" - mac_vendor_lookup = AsyncMacLookup() - try: - await mac_vendor_lookup.update_vendors() - except: - try: - await mac_vendor_lookup.load_vendors() - except: - pass + """Set up device tracker for OPNsense component.""" dev_reg = async_get_dev_reg(hass) - @callback - def process_entities_callback( - hass: HomeAssistant, config_entry: ConfigEntry - ) -> list[OPNsenseScannerEntity]: - # options = config_entry.options - data = hass.data[DOMAIN][config_entry.entry_id] - previous_mac_addresses = config_entry.data.get(TRACKED_MACS, []) - coordinator = data[DEVICE_TRACKER_COORDINATOR] - state = coordinator.data - # seems unlikely *all* devices are intended to be monitored - # disable by default and let users enable specific entries they care about - enabled_default = False - device_per_arp_entry = False - - entities = [] - mac_addresses = [] - - # use configured mac addresses if setup, otherwise create an entity per arp entry - configured_mac_addresses = config_entry.options.get(CONF_DEVICES, []) - if configured_mac_addresses: - mac_addresses = configured_mac_addresses - enabled_default = True - else: - if device_per_arp_entry: - arp_entries = dict_get(state, "arp_table") - if not arp_entries: - return [] - - mac_addresses: list = [ - mac_address.lower() - for arp_entry in arp_entries - if (mac_address := arp_entry.get("mac", "")) - ] - + data = hass.data[DOMAIN][config_entry.entry_id] + previous_mac_addresses: list = config_entry.data.get(TRACKED_MACS, []) + coordinator: OPNsenseDataUpdateCoordinator = data[DEVICE_TRACKER_COORDINATOR] + state = coordinator.data + + enabled_default = False + entities: list = [] + mac_addresses: list = [] + + # use configured mac addresses if setup, otherwise create an entity per arp entry + arp_entries: list = dict_get(state, "arp_table") + if not isinstance(arp_entries, list): + arp_entries = [] + devices: list = [] + mac_addresses = [] + configured_mac_addresses = config_entry.options.get(CONF_DEVICES, []) + if configured_mac_addresses and config_entry.options.get( + CONF_DEVICE_TRACKER_ENABLED, DEFAULT_DEVICE_TRACKER_ENABLED + ): + _LOGGER.debug( + f"[device_tracker async_setup_entry] configured_mac_addresses: {configured_mac_addresses}" + ) + enabled_default = True + mac_addresses = configured_mac_addresses for mac_address in mac_addresses: - mac_vendor = None - try: - mac_vendor = lookup_mac(mac_vendor_lookup, mac_address) - except: - pass - - entity = OPNsenseScannerEntity( - hass, - config_entry, - coordinator, - enabled_default, - mac_address, - mac_vendor, - ) - - entities.append(entity) - - # Get the MACs that need to be removed and remove their devices - for mac_address in list(set(previous_mac_addresses) - set(mac_addresses)): - device = dev_reg.async_get_device( - {}, {(CONNECTION_NETWORK_MAC, mac_address)} - ) - if device: - dev_reg.async_remove_device(device.id) + device: Mapping[str, Any] = {"mac": mac_address} + for arp_entry in arp_entries: + if mac_address == arp_entry.get("mac", ""): + for attr in ["hostname", "manufacturer"]: + try: + if arp_entry.get(attr, None): + device.update({attr: arp_entry.get(attr, None)}) + except (TypeError, KeyError, AttributeError): + pass + devices.append(device) + elif config_entry.options.get( + CONF_DEVICE_TRACKER_ENABLED, DEFAULT_DEVICE_TRACKER_ENABLED + ): + for arp_entry in arp_entries: + mac_address = arp_entry.get("mac", None) + if mac_address and mac_address not in mac_addresses: + device: Mapping[str, Any] = {"mac": mac_address} + for attr in ["hostname", "manufacturer"]: + try: + if arp_entry.get(attr, None): + device.update({attr: arp_entry.get(attr, None)}) + except (TypeError, KeyError, AttributeError): + pass + mac_addresses.append(mac_address) + devices.append(device) + + for device in devices: + entity = OPNsenseScannerEntity( + config_entry=config_entry, + coordinator=coordinator, + enabled_default=enabled_default, + mac=device.get("mac", None), + mac_vendor=device.get("manufacturer", None), + hostname=device.get("hostname", None), + ) + entities.append(entity) - if set(mac_addresses) != set(previous_mac_addresses): - data[SHOULD_RELOAD] = False - new_data = config_entry.data.copy() - new_data[TRACKED_MACS] = mac_addresses.copy() - hass.config_entries.async_update_entry(config_entry, data=new_data) + # Get the MACs that need to be removed and remove their devices + for mac_address in list(set(previous_mac_addresses) - set(mac_addresses)): + device = dev_reg.async_get_device({}, {(CONNECTION_NETWORK_MAC, mac_address)}) + if device: + dev_reg.async_remove_device(device.id) - return entities + if set(mac_addresses) != set(previous_mac_addresses): + data[SHOULD_RELOAD] = False + new_data = config_entry.data.copy() + new_data[TRACKED_MACS] = mac_addresses.copy() + hass.config_entries.async_update_entry(config_entry, data=new_data) - cem = CoordinatorEntityManager( - hass, - hass.data[DOMAIN][config_entry.entry_id][DEVICE_TRACKER_COORDINATOR], - config_entry, - process_entities_callback, - async_add_entities, - ) - cem.process_entities() + _LOGGER.debug(f"[device_tracker async_setup_entry] entities: {len(entities)}") + async_add_entities(entities) -class OPNsenseScannerEntity(OPNsenseEntity, ScannerEntity): +class OPNsenseScannerEntity(OPNsenseEntity, ScannerEntity, RestoreEntity): """Represent a scanned device.""" def __init__( self, - hass: HomeAssistant, config_entry: ConfigEntry, coordinator: OPNsenseDataUpdateCoordinator, enabled_default: bool, mac: str, - mac_vendor: str, + mac_vendor: str | None, + hostname: str | None, ) -> None: """Set up the OPNsense scanner entity.""" super().__init__(config_entry, coordinator, unique_id_suffix=f"mac_{mac}") - self._mac_address = mac - self._mac_vendor = mac_vendor - self._last_known_ip = None - self._last_known_hostname = None - self._last_known_connected_time = None - self._extra_state = {} - - self._attr_entity_registry_enabled_default = enabled_default - - def _get_opnsense_arp_entry(self) -> dict[str, str]: - state = self.coordinator.data - arp_table = dict_get(state, "arp_table") - if arp_table is None: - return None - for entry in arp_table: - if entry.get("mac", "").lower() == self._mac_address: - return entry - - return None + self._mac_vendor: str | None = mac_vendor + self._attr_name: str | None = f"{self.opnsense_device_name} {hostname or mac}" + self._last_known_ip: str | None = None + self._last_known_hostname: str | None = None + self._is_connected: bool = False + self._last_known_connected_time: str | None = None + self._attr_entity_registry_enabled_default: bool = enabled_default + self._attr_extra_state_attributes: Mapping[str, Any] = {} + self._attr_hostname: str | None = hostname + self._attr_ip_address: str | None = None + self._attr_mac_address: str | None = mac + self._attr_source_type: SourceType = SourceType.ROUTER + self._available: bool = ( + False # Move this to OPNsenseEntity once all entity-types are updated + ) + # Move this to OPNsenseEntity once all entity-types are updated @property def available(self) -> bool: - state = self.coordinator.data - arp_table = dict_get(state, "arp_table") - if arp_table is None: - return False - return super().available - - @property - def source_type(self) -> str: - """Return the source type, eg gps or router, of the device.""" - return SourceType.ROUTER + return self._available @property - def extra_state_attributes(self) -> Mapping[str, Any] | None: - """Return extra state attributes.""" - return self._extra_state_attributes + def source_type(self) -> SourceType: + return self._attr_source_type @property - def _extra_state_attributes(self) -> Mapping[str, Any] | None: - """Return extra state attributes.""" - entry = self._get_opnsense_arp_entry() - if entry is not None: - for prop_name in ["interface", "expires", "type"]: - self._extra_state[prop_name] = entry.get(prop_name) - - if self._last_known_hostname is not None: - self._extra_state["last_known_hostname"] = self._last_known_hostname - - if self._last_known_ip is not None: - self._extra_state["last_known_ip"] = self._last_known_ip - - if self._last_known_connected_time is not None: - self._extra_state["last_known_connected_time"] = ( - self._last_known_connected_time - ) - - return self._extra_state + def is_connected(self) -> bool: + return self._is_connected @property def ip_address(self) -> str | None: - """Return the primary ip address of the device.""" - return self._ip_address - - @property - def _ip_address(self) -> str | None: - """Return the primary ip address of the device.""" - entry = self._get_opnsense_arp_entry() - if entry is None: - return None - - ip_address = entry.get("ip") - if ip_address is not None and len(ip_address) > 0: - self._last_known_ip = ip_address - return ip_address + return self._attr_ip_address @property def mac_address(self) -> str | None: - """Return the mac address of the device.""" - return self._mac_address + return self._attr_mac_address @property def hostname(self) -> str | None: - """Return hostname of the device.""" - return self._hostname + return self._attr_hostname @property - def _hostname(self) -> str | None: - """Return hostname of the device.""" - entry = self._get_opnsense_arp_entry() - if entry is None: - return None - value = entry.get("hostname").strip("?") - if len(value) > 0: - self._last_known_hostname = value - return value - return None + def unique_id(self) -> str | None: + return self._attr_unique_id @property - def name(self) -> str: - """Return the name of the device.""" - identifier = self.hostname or self._last_known_hostname or self._mac_address - return f"{self.opnsense_device_name} {identifier}" + def entity_registry_enabled_default(self) -> bool: + return self._attr_entity_registry_enabled_default - @property - def device_info(self) -> DeviceInfo: - """Return the device info.""" - return DeviceInfo( - connections={(CONNECTION_NETWORK_MAC, self.mac_address)}, - default_manufacturer=self._mac_vendor, - default_name=self.name, - via_device=(DOMAIN, self.opnsense_device_unique_id), - ) + @callback + def _handle_coordinator_update(self) -> None: + state = self.coordinator.data + arp_table = dict_get(state, "arp_table") + if not isinstance(arp_table, list): + self._available = False + return + self._available = True + entry: Mapping[str:Any] | None = None + for arp_entry in arp_table: + if arp_entry.get("mac", "").lower() == self._attr_mac_address: + entry = arp_entry + break + + # _LOGGER.debug(f"[OPNsenseScannerEntity handle_coordinator_update] entry: {entry}") + try: + self._attr_ip_address = ( + entry.get("ip") if len(entry.get("ip")) > 0 else None + ) + except (TypeError, KeyError, AttributeError): + self._attr_ip_address = None + + if self._attr_ip_address: + self._last_known_ip = self._attr_ip_address - @property - def icon(self) -> str: - """Return device icon.""" try: - return "mdi:lan-connect" if self.is_connected else "mdi:lan-disconnect" - except: - return "mdi:lan-disconnect" + self._attr_hostname = ( + entry.get("hostname").strip("?") + if len(entry.get("hostname").strip("?")) > 0 + else None + ) + except (TypeError, KeyError, AttributeError): + self._attr_hostname = None + + if self._attr_hostname: + self._last_known_hostname = self._attr_hostname - @property - def is_connected(self) -> bool: - """Return true if the device is connected to the network.""" - state = self.coordinator.data update_time = state["update_time"] - entry = self._get_opnsense_arp_entry() if entry is None: - if self._last_known_ip is not None and len(self._last_known_ip) > 0: + if self._last_known_ip: # force a ping to _last_known_ip to possibly recreate arp entry? pass @@ -299,46 +238,114 @@ def is_connected(self) -> bool: current_time = int(time.time()) elapsed = current_time - self._last_known_connected_time if elapsed < device_tracker_consider_home: - return True + self._is_connected = True - return False - # TODO: check "expires" here to add more honed in logic? - # TODO: clear cache under certain scenarios? - ip_address = entry.get("ip") - if ip_address is not None and len(ip_address) > 0: - client = self._get_opnsense_client() - self.hass.add_job(client.delete_arp_entry, ip_address) + self._is_connected = False + else: + # TODO: check "expires" here to add more honed in logic? + # TODO: clear cache under certain scenarios? - self._last_known_connected_time = int(update_time) + # Why was this being done? Remove it? + # ip_address = entry.get("ip") + # if ip_address is not None and len(ip_address) > 0: + # self.hass.add_job(self._client.delete_arp_entry, ip_address) - return True + self._last_known_connected_time = datetime.fromtimestamp( + int(update_time), + tz=datetime.now().astimezone().tzinfo, + ) + self._is_connected = True + + ha_to_opnsense: Mapping[str, Any] = { + "interface": "intf_description", + "expires": "expires", + "type": "type", + } + for prop_name in ["interface", "expires", "type"]: + try: + prop = entry.get(ha_to_opnsense[prop_name]) + if prop: + if prop_name == "expires": + if prop == -1: + self._attr_extra_state_attributes[prop_name] = "Never" + else: + self._attr_extra_state_attributes[prop_name] = ( + datetime.now().astimezone() + timedelta(seconds=prop) + ) + else: + self._attr_extra_state_attributes[prop_name] = prop + except (TypeError, KeyError, AttributeError): + pass - async def async_added_to_hass(self) -> None: - """Handle entity which will be added.""" - await super().async_added_to_hass() - state = await self.async_get_last_state() - if state is None: - return + if self._attr_hostname is None and self._last_known_hostname: + self._attr_extra_state_attributes["last_known_hostname"] = ( + self._last_known_hostname + ) + else: + self._attr_extra_state_attributes.pop("last_known_hostname", None) - if state.attributes is None: + if self._attr_ip_address is None and self._last_known_ip: + self._attr_extra_state_attributes["last_known_ip"] = self._last_known_ip + else: + self._attr_extra_state_attributes.pop("last_known_ip", None) + + if self._last_known_connected_time is not None: + self._attr_extra_state_attributes["last_known_connected_time"] = ( + self._last_known_connected_time + ) + + try: + self._attr_icon = ( + "mdi:lan-connect" if self.is_connected else "mdi:lan-disconnect" + ) + except (TypeError, KeyError, AttributeError): + self._attr_icon = "mdi:lan-disconnect" + + self.async_write_ha_state() + _LOGGER.debug( + f"[OPNsenseScannerEntity handle_coordinator_update] Name: {self.name}, " + f"unique_id: {self.unique_id}, attr_unique_id: {self._attr_unique_id}, " + f"available: {self.available}, is_connected: {self.is_connected}, " + f"hostname: {self.hostname}, ip_address: {self.ip_address}, " + f"last_known_hostname: {self._last_known_hostname}, last_known_ip: {self._last_known_ip}, " + f"last_known_connected_time: {self._last_known_connected_time}, icon: {self.icon}, " + f"extra_state_atrributes: {self.extra_state_attributes}" + ) + + @property + def device_info(self) -> DeviceInfo: + """Return the device info.""" + return DeviceInfo( + connections={(CONNECTION_NETWORK_MAC, self.mac_address)}, + default_manufacturer=self._mac_vendor, + default_name=self.name, + via_device=(DOMAIN, self.opnsense_device_unique_id), + ) + + async def _restore_last_state(self) -> None: + state = await self.async_get_last_state() + if state is None or state.attributes is None: return state = state.attributes + + self._last_known_hostname = state.get("last_known_hostname", None) + self._last_known_ip = state.get("last_known_ip", None) + for attr in [ "interface", "expires", "type", - "last_known_ip", - "last_known_hostname", + "last_known_connected_time", ]: - value = state.get(attr, None) - if value is not None: - self._extra_state[attr] = value - if attr == "last_known_hostname": - self._last_known_hostname = value - - if attr == "last_known_ip": - self._last_known_ip = value + try: + value = state.get(attr, None) + if value: + self._attr_extra_state_attributes[attr] = value + except (TypeError, KeyError, AttributeError): + pass - if attr == "last_known_connected_time": - self._last_known_connected_time = value + async def async_added_to_hass(self) -> None: + await super().async_added_to_hass() + await self._restore_last_state() + self._handle_coordinator_update() diff --git a/custom_components/opnsense/manifest.json b/custom_components/opnsense/manifest.json index e54feba..1b5c320 100644 --- a/custom_components/opnsense/manifest.json +++ b/custom_components/opnsense/manifest.json @@ -7,6 +7,6 @@ "documentation": "https://github.com/travisghansen/hass-opnsense", "iot_class": "local_polling", "issue_tracker": "https://github.com/travisghansen/hass-opnsense/issues", - "requirements": ["mac-vendor-lookup>=0.1.11", "python-dateutil", "awesomeversion"], + "requirements": ["python-dateutil", "awesomeversion"], "version":"v0.3.2" } diff --git a/requirements_dev.txt b/requirements_dev.txt index f4e0963..f825700 100644 --- a/requirements_dev.txt +++ b/requirements_dev.txt @@ -1,6 +1,4 @@ black homeassistant isort -mac_vendor_lookup python-dateutil -async-timeout