diff --git a/custom_components/mikrotik_router/device_tracker.py b/custom_components/mikrotik_router/device_tracker.py index 18a18ff..b8f1eb7 100644 --- a/custom_components/mikrotik_router/device_tracker.py +++ b/custom_components/mikrotik_router/device_tracker.py @@ -1,6 +1,7 @@ """Support for the Mikrotik Router device tracker.""" import logging +from datetime import timedelta from homeassistant.components.device_tracker.config_entry import ScannerEntity from homeassistant.components.device_tracker.const import SOURCE_TYPE_ROUTER @@ -11,12 +12,16 @@ from homeassistant.const import ( from homeassistant.core import callback from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC from homeassistant.helpers.dispatcher import async_dispatcher_connect -from homeassistant.util.dt import get_age +from homeassistant.util.dt import get_age, utcnow from .const import ( DOMAIN, DATA_CLIENT, ATTRIBUTION, + CONF_TRACK_HOSTS, + DEFAULT_TRACK_HOSTS, + CONF_TRACK_HOSTS_TIMEOUT, + DEFAULT_TRACK_HOST_TIMEOUT, ) _LOGGER = logging.getLogger(__name__) @@ -42,6 +47,7 @@ DEVICE_ATTRIBUTES_HOST = [ "address", "mac-address", "interface", + "source", "last-seen", ] @@ -58,6 +64,18 @@ def format_attribute(attr): return res +# --------------------------- +# format_value +# --------------------------- +def format_value(res): + res = res.replace("dhcp", "DHCP") + res = res.replace("dns", "DNS") + res = res.replace("capsman", "CAPsMAN") + res = res.replace("wireless", "Wireless") + res = res.replace("restored", "Restored") + return res + + # --------------------------- # async_setup_entry # --------------------------- @@ -70,7 +88,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): @callback def update_controller(): """Update the values of the controller.""" - update_items(inst, mikrotik_controller, async_add_entities, tracked) + update_items(inst, config_entry, mikrotik_controller, async_add_entities, tracked) mikrotik_controller.listeners.append( async_dispatcher_connect( @@ -85,7 +103,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): # update_items # --------------------------- @callback -def update_items(inst, mikrotik_controller, async_add_entities, tracked): +def update_items(inst, config_entry, mikrotik_controller, async_add_entities, tracked): """Update tracked device state from the controller.""" new_tracked = [] @@ -106,7 +124,7 @@ def update_items(inst, mikrotik_controller, async_add_entities, tracked): tracked[item_id].async_schedule_update_ha_state() continue - tracked[item_id] = sid_func(inst, uid, mikrotik_controller) + tracked[item_id] = sid_func(inst, uid, mikrotik_controller, config_entry) new_tracked.append(tracked[item_id]) if new_tracked: @@ -119,7 +137,7 @@ def update_items(inst, mikrotik_controller, async_add_entities, tracked): class MikrotikControllerPortDeviceTracker(ScannerEntity): """Representation of a network port.""" - def __init__(self, inst, uid, mikrotik_controller): + def __init__(self, inst, uid, mikrotik_controller, _): """Set up tracked port.""" self._inst = inst self._ctrl = mikrotik_controller @@ -200,7 +218,6 @@ class MikrotikControllerPortDeviceTracker(ScannerEntity): def device_state_attributes(self): """Return the port state attributes.""" attributes = self._attrs - for variable in DEVICE_ATTRIBUTES_IFACE: if variable in self._data: attributes[format_attribute(variable)] = self._data[variable] @@ -214,16 +231,30 @@ class MikrotikControllerPortDeviceTracker(ScannerEntity): class MikrotikControllerHostDeviceTracker(ScannerEntity): """Representation of a network device.""" - def __init__(self, inst, uid, mikrotik_controller): + def __init__(self, inst, uid, mikrotik_controller, config_entry): """Set up tracked port.""" self._inst = inst self._ctrl = mikrotik_controller self._data = mikrotik_controller.data["host"][uid] + self._config_entry = config_entry self._attrs = { ATTR_ATTRIBUTION: ATTRIBUTION, } + @property + def option_track_network_hosts(self): + """Config entry option to not track ARP.""" + return self._config_entry.options.get(CONF_TRACK_HOSTS, DEFAULT_TRACK_HOSTS) + + @property + def option_track_network_hosts_timeout(self): + """Config entry option scan interval.""" + track_network_hosts_timeout = self._config_entry.options.get( + CONF_TRACK_HOSTS_TIMEOUT, DEFAULT_TRACK_HOST_TIMEOUT + ) + return timedelta(seconds=track_network_hosts_timeout) + @property def entity_registry_enabled_default(self): """Return if the entity should be enabled when first added to the entity registry.""" @@ -244,7 +275,16 @@ class MikrotikControllerHostDeviceTracker(ScannerEntity): @property def is_connected(self): """Return true if the host is connected to the network.""" - return self._data["available"] + if self._data["source"] in ["capsman", "wireless"]: + return self._data["available"] + + if ( + self._data["last-seen"] + and (utcnow() - self._data["last-seen"]) + < self.option_track_network_hosts_timeout + ): + return True + return False @property def source_type(self): @@ -264,17 +304,27 @@ class MikrotikControllerHostDeviceTracker(ScannerEntity): @property def available(self) -> bool: """Return if controller is available.""" + if not self.option_track_network_hosts: + return False + return self._ctrl.connected() @property def icon(self): """Return the icon.""" - if self._data["available"]: - icon = "mdi:lan-connect" - else: - icon = "mdi:lan-disconnect" + if self._data["source"] in ["capsman", "wireless"]: + if self._data["available"]: + return "mdi:lan-connect" + else: + return "mdi:lan-disconnect" - return icon + if ( + self._data["last-seen"] + and (utcnow() - self._data["last-seen"]) + < self.option_track_network_hosts_timeout + ): + return "mdi:lan-connect" + return "mdi:lan-disconnect" @property def device_info(self): @@ -299,7 +349,6 @@ class MikrotikControllerHostDeviceTracker(ScannerEntity): def device_state_attributes(self): """Return the host state attributes.""" attributes = self._attrs - for variable in DEVICE_ATTRIBUTES_HOST: if variable in self._data: if variable == "last-seen": @@ -308,6 +357,9 @@ class MikrotikControllerHostDeviceTracker(ScannerEntity): else: attributes[format_attribute(variable)] = "unknown" else: - attributes[format_attribute(variable)] = self._data[variable] + if self._data[variable] in ["dhcp", "dns", "capsman", "wireless", "restored"]: + attributes[format_attribute(variable)] = format_value(self._data[variable]) + else: + attributes[format_attribute(variable)] = self._data[variable] return attributes