Merge branch 'master' of https://github.com/ivanpavlina/homeassistant-mikrotik_router into pr/2

custom_components/mikrotik_router/__init__.py

@@ -54,6 +54,7 @@ async def async_setup_entry(hass, config_entry):
         traffic_type
     )
     await mikrotik_controller.hwinfo_update()
+
     await mikrotik_controller.async_update()
 
     if not mikrotik_controller.data:

custom_components/mikrotik_router/config_flow.py

@@ -81,7 +81,7 @@ class MikrotikControllerConfigFlow(ConfigFlow, domain=DOMAIN):
                 username=user_input["username"],
                 password=user_input["password"],
                 port=user_input["port"],
-                use_ssl=user_input["ssl"],
+                use_ssl=user_input["ssl"]
             )
             if not api.connect():
                 errors[CONF_HOST] = api.error

custom_components/mikrotik_router/mikrotik_controller.py

@@ -66,13 +66,17 @@ class MikrotikControllerData:
             "dhcp-network": {},
             "dhcp": {},
             "host": {},
+            "accounting": {}
         }
 
+        self.local_dhcp_networks = []
+
         self.listeners = []
         self.lock = asyncio.Lock()
 
         self.api = MikrotikAPI(host, username, password, port, use_ssl)
 
+        self.raw_arp_entries = []
         self.nat_removed = {}
 
         async_track_time_interval(
@@ -205,6 +209,8 @@ class MikrotikControllerData:
         await self.hass.async_add_executor_job(self.get_system_resource)
         await self.hass.async_add_executor_job(self.get_script)
         await self.hass.async_add_executor_job(self.get_queue)
+        await self.hass.async_add_executor_job(self.get_dhcp)
+        await self.hass.async_add_executor_job(self.get_accounting)
 
         async_dispatcher_send(self.hass, self.signal_update)
 
@@ -348,6 +354,7 @@ class MikrotikControllerData:
     def update_arp(self, mac2ip, bridge_used):
         """Get list of hosts in ARP for interface client data from Mikrotik"""
         data = self.api.path("/ip/arp")
+        self.raw_arp_entries = data
         if not data:
             return mac2ip, bridge_used
 
@@ -561,7 +568,7 @@ class MikrotikControllerData:
             self.data["resource"]["hdd-usage"] = "unknown"
 
     # ---------------------------
-    #   get_system_routerboard
+    #   get_firmware_update
     # ---------------------------
     def get_firmware_update(self):
         """Check for firmware update on Mikrotik"""
@@ -736,6 +743,20 @@ class MikrotikControllerData:
             ]
         )
 
+        # Build list of local DHCP networks
+        dhcp_networks = parse_api(
+            data={},
+            source=self.api.path("/ip/dhcp-server/network"),
+            key="address",
+            vals=[
+                {"name": "address"},
+            ],
+            ensure_vals=[
+                {"name": "address"},
+            ]
+        )
+        self.local_dhcp_networks = [IPv4Network(network) for network in dhcp_networks]
+
         self.data["dhcp"] = parse_api(
             data=self.data["dhcp"],
             source=self.api.path("/ip/dhcp-server/lease"),
@@ -747,7 +768,7 @@ class MikrotikControllerData:
                 {"name": "status", "default": "unknown"},
                 {"name": "last-seen", "default": "unknown"},
                 {"name": "server", "default": "unknown"},
-                {"name": "comment", "default": ""},
+                {"name": "comment"},
             ],
             ensure_vals=[
                 {"name": "interface"},
@@ -827,3 +848,178 @@ class MikrotikControllerData:
             # Update last seen
             if self.data["host"][uid]["available"]:
                 self.data["host"][uid]["last-seen"] = utcnow()
+            self.data["dhcp"][uid]['interface'] = \
+                self.data["dhcp-server"][self.data["dhcp"][uid]['server']]["interface"]
+
+            self.data["dhcp"][uid]['available'] = \
+                self.api.arp_ping(self.data["dhcp"][uid]['address'], self.data["dhcp"][uid]['interface'])
+
+    def _address_part_of_local_network(self, address):
+        address = ip_address(address)
+        for network in self.local_dhcp_networks:
+            if address in network:
+                return True
+        return False
+
+    def _get_accounting_mac_by_ip(self, requested_ip):
+        for mac, vals in self.data['accounting'].items():
+            if vals.get('address') is requested_ip:
+                return mac
+        return None
+
+    def get_accounting(self):
+        """Get Accounting data from Mikrotik"""
+        # Check if accounting and account-local-traffic is enabled
+        accounting_enabled, local_traffic_enabled = self.api.is_accounting_and_local_traffic_enabled()
+        traffic_type, traffic_div = self._get_traffic_type_and_div()
+
+        if not accounting_enabled:
+            # If any hosts were created return counters to 0 so sensors wont get stuck on last value
+            for mac, vals in self.data["accounting"].items():
+                self.data['accounting'][mac]["tx-rx-attr"] = traffic_type
+                if 'wan-tx' in vals:
+                    self.data["accounting"][mac]['wan-tx'] = 0.0
+                if 'wan-rx' in vals:
+                    self.data["accounting"][mac]['wan-rx'] = 0.0
+                if 'lan-tx' in vals:
+                    self.data["accounting"][mac]['lan-tx'] = 0.0
+                if 'lan-rx' in vals:
+                    self.data["accounting"][mac]['lan-rx'] = 0.0
+            return
+
+        # Build missing hosts from already retrieved DHCP Server leases
+        for mac, vals in self.data["dhcp"].items():
+            if mac not in self.data["accounting"]:
+                self.data["accounting"][mac] = {
+                    'address': vals['address'],
+                    'mac-address': vals['mac-address'],
+                    'host-name': vals['host-name'],
+                    'comment': vals['comment']
+                }
+
+        # Build missing hosts from already retrieved ARP list
+        host_update_from_arp = False
+        for entry in self.raw_arp_entries:
+            if entry['mac-address'] not in self.data["accounting"]:
+                self.data["accounting"][entry['mac-address']] = {
+                    'address': entry['address'],
+                    'mac-address': entry['mac-address'],
+                    'host-name': '',
+                    'comment': ''
+                }
+                host_update_from_arp = True
+
+        # If some host was added from ARP table build new host-name for it from static DNS entry. Fallback to MAC
+        if host_update_from_arp:
+            dns_data = parse_api(
+                data={},
+                source=self.api.path("/ip/dns/static"),
+                key="address",
+                vals=[
+                    {"name": "address"},
+                    {"name": "name"},
+                ],
+            )
+
+            # Try to build hostname from DNS static entry
+            for mac, vals in self.data["accounting"].items():
+                if not str(vals.get('host-name', '')).strip() or vals['host-name'] is 'unknown':
+                    if vals['address'] in dns_data and str(dns_data[vals['address']].get('name', '')).strip():
+                        self.data["accounting"][mac]['host-name'] = dns_data[vals['address']]['name']
+
+        # Check if any host still have empty 'host-name'. Default it to MAC.
+        # Same check for 'comment' (pretty name)
+        for mac, vals in self.data["accounting"].items():
+            if not str(vals.get('host-name', '')).strip() or vals['host-name'] is 'unknown':
+                self.data["accounting"][mac]['host-name'] = mac
+            if not str(vals.get('comment', '')).strip() or vals['host-name'] is 'comment':
+                self.data["accounting"][mac]['comment'] = mac
+
+        _LOGGER.debug(f"Working with {len(self.data['accounting'])} accounting devices")
+
+        # Build temp accounting values dict with ip address as key
+        # Also set traffic type for each item
+        tmp_accounting_values = {}
+        for mac, vals in self.data['accounting'].items():
+            tmp_accounting_values[vals['address']] = {
+                "wan-tx": 0,
+                "wan-rx": 0,
+                "lan-tx": 0,
+                "lan-rx": 0
+            }
+            self.data['accounting'][mac]["tx-rx-attr"] = traffic_type
+
+        time_diff = self.api.take_accounting_snapshot()
+        if time_diff:
+            accounting_data = parse_api(
+                data={},
+                source=self.api.path("/ip/accounting/snapshot"),
+                key=".id",
+                vals=[
+                    {"name": ".id"},
+                    {"name": "src-address"},
+                    {"name": "dst-address"},
+                    {"name": "bytes", "default": 0},
+                ],
+            )
+
+            for item in accounting_data.values():
+                source_ip = str(item.get('src-address')).strip()
+                destination_ip = str(item.get('dst-address')).strip()
+                bits_count = int(str(item.get('bytes')).strip()) * 8
+
+                if self._address_part_of_local_network(source_ip) and self._address_part_of_local_network(destination_ip):
+                    # LAN TX/RX
+                    if source_ip in tmp_accounting_values:
+                        tmp_accounting_values[source_ip]['lan-tx'] += bits_count
+                    if destination_ip in tmp_accounting_values:
+                        tmp_accounting_values[destination_ip]['lan-rx'] += bits_count
+                elif self._address_part_of_local_network(source_ip) and \
+                        not self._address_part_of_local_network(destination_ip):
+                    # WAN TX
+                    if source_ip in tmp_accounting_values:
+                        tmp_accounting_values[source_ip]['wan-tx'] += bits_count
+                elif not self._address_part_of_local_network(source_ip) and \
+                        self._address_part_of_local_network(destination_ip):
+                    # WAN RX
+                    if destination_ip in tmp_accounting_values:
+                        tmp_accounting_values[destination_ip]['wan-rx'] += bits_count
+
+            # Now that we have sum of all traffic in bytes for given period
+            #   calculate real throughput and transform it to appropriate unit
+            for addr in tmp_accounting_values:
+                mac = self._get_accounting_mac_by_ip(addr)
+                if not mac:
+                    _LOGGER.debug(f"Address {addr} not found in accounting data, skipping update")
+                    continue
+
+                self.data['accounting'][mac]['wan-tx'] = round(
+                    tmp_accounting_values[addr]['wan-tx'] / time_diff * traffic_div, 2)
+                self.data['accounting'][mac]['wan-rx'] = round(
+                    tmp_accounting_values[addr]['wan-rx'] / time_diff * traffic_div, 2)
+
+                if local_traffic_enabled:
+                    self.data['accounting'][mac]['lan-tx'] = round(
+                        tmp_accounting_values[addr]['lan-tx'] / time_diff * traffic_div, 2)
+                    self.data['accounting'][mac]['lan-rx'] = round(
+                        tmp_accounting_values[addr]['lan-rx'] / time_diff * traffic_div, 2)
+                else:
+                    # If local traffic was enabled earlier and then disabled return counters for LAN traffic to 0
+                    if 'lan-tx' in self.data['accounting'][mac]:
+                        self.data['accounting'][mac]['lan-tx'] = 0.0
+                    if 'lan-rx' in self.data['accounting'][mac]:
+                        self.data['accounting'][mac]['lan-rx'] = 0.0
+        else:
+            # No time diff, just initialize/return counters to 0 for all
+            for addr in tmp_accounting_values:
+                mac = self._get_accounting_mac_by_ip(addr)
+                if not mac:
+                    _LOGGER.debug(f"Address {addr} not found in accounting data, skipping update")
+                    continue
+
+                self.data['accounting'][mac]['wan-tx'] = 0.0
+                self.data['accounting'][mac]['wan-rx'] = 0.0
+
+                if local_traffic_enabled:
+                    self.data['accounting'][mac]['lan-tx'] = 0.0
+                    self.data['accounting'][mac]['lan-rx'] = 0.0

custom_components/mikrotik_router/mikrotikapi.py

@@ -60,6 +60,7 @@ class MikrotikAPI:
         self._connection_retry_sec = 58
         self.error = None
         self.connection_error_reported = False
+        self.accounting_last_run = None
 
         # Default ports
         if not self._port:
@@ -427,6 +428,7 @@ class MikrotikAPI:
             self.disconnect()
             self.lock.release()
             return False
+
         except (
             librouteros_custom.exceptions.TrapError,
             librouteros_custom.exceptions.MultiTrapError,
@@ -463,3 +465,95 @@ class MikrotikAPI:
                 return True
 
         return False
+
+    @staticmethod
+    def _current_milliseconds():
+        from time import time
+        return int(round(time() * 1000))
+
+    def is_accounting_and_local_traffic_enabled(self) -> (bool, bool):
+        # Returns:
+        #   1st bool: Is accounting enabled
+        #   2nd bool: Is account-local-traffic enabled
+
+        if not self.connection_check():
+            return False, False
+
+        response = self.path("/ip/accounting")
+        if response is None:
+            return False, False
+
+        for item in response:
+            if 'enabled' not in item:
+                continue
+            if not item['enabled']:
+                return False, False
+
+        for item in response:
+            if 'account-local-traffic' not in item:
+                continue
+            if not item['account-local-traffic']:
+                return True, False
+
+        return True, True
+
+    # ---------------------------
+    #   take_accounting_snapshot
+    #   Returns float -> period in seconds between last and current run
+    # ---------------------------
+    def take_accounting_snapshot(self) -> float:
+        """Get accounting data"""
+        if not self.connection_check():
+            return 0
+
+        accounting = self.path("/ip/accounting", return_list=False)
+
+        self.lock.acquire()
+        try:
+            # Prepare command
+            take = accounting('snapshot/take')
+        except librouteros_custom.exceptions.ConnectionClosed:
+            self.disconnect()
+            self.lock.release()
+            return 0
+        except (
+                librouteros_custom.exceptions.TrapError,
+                librouteros_custom.exceptions.MultiTrapError,
+                librouteros_custom.exceptions.ProtocolError,
+                librouteros_custom.exceptions.FatalError,
+                ssl.SSLError,
+                BrokenPipeError,
+                OSError,
+                ValueError,
+        ) as api_error:
+            self.disconnect("accounting_snapshot", api_error)
+            self.lock.release()
+            return 0
+        except:
+            self.disconnect("accounting_snapshot")
+            self.lock.release()
+            return 0
+
+        try:
+            list(take)
+        except librouteros_custom.exceptions.ConnectionClosed as api_error:
+            self.disconnect("accounting_snapshot", api_error)
+            self.lock.release()
+            return 0
+        except:
+            self.disconnect("accounting_snapshot")
+            self.lock.release()
+            return 0
+
+        self.lock.release()
+
+        # First request will be discarded because we cannot know when the last data was retrieved
+        # prevents spikes in data
+        if not self.accounting_last_run:
+            self.accounting_last_run = self._current_milliseconds()
+            return 0
+
+        # Calculate time difference in seconds and return
+        time_diff = self._current_milliseconds() - self.accounting_last_run
+        self.accounting_last_run = self._current_milliseconds()
+        return time_diff / 1000

custom_components/mikrotik_router/sensor.py

@@ -12,6 +12,19 @@ from .const import (DOMAIN, DATA_CLIENT, ATTRIBUTION)
 
 _LOGGER = logging.getLogger(__name__)
 
+
+# ---------------------------
+#   format_attribute
+# ---------------------------
+def format_attribute(attr):
+    res = attr.replace("-", " ")
+    res = res.capitalize()
+    res = res.replace(" ip ", " IP ")
+    res = res.replace(" mac ", " MAC ")
+    res = res.replace(" mtu", " MTU")
+    return res
+
+
 ATTR_ICON = "icon"
 ATTR_LABEL = "label"
 ATTR_UNIT = "unit"
@@ -66,8 +79,54 @@ SENSOR_TYPES = {
         ATTR_PATH: "interface",
         ATTR_ATTR: "rx-bits-per-second",
     },
+    "accounting_lan_tx": {
+        ATTR_DEVICE_CLASS: None,
+        ATTR_ICON: "mdi:download-network",
+        ATTR_LABEL: "LAN TX",
+        ATTR_GROUP: "Accounting",
+        ATTR_UNIT: "ps",
+        ATTR_UNIT_ATTR: "tx-rx-attr",
+        ATTR_PATH: "accounting",
+        ATTR_ATTR: "lan-tx",
+    },
+    "accounting_lan_rx": {
+        ATTR_DEVICE_CLASS: None,
+        ATTR_ICON: "mdi:upload-network",
+        ATTR_LABEL: "LAN RX",
+        ATTR_GROUP: "Accounting",
+        ATTR_UNIT: "ps",
+        ATTR_UNIT_ATTR: "tx-rx-attr",
+        ATTR_PATH: "accounting",
+        ATTR_ATTR: "lan-rx",
+    },
+    "accounting_wan_tx": {
+        ATTR_DEVICE_CLASS: None,
+        ATTR_ICON: "mdi:download-network",
+        ATTR_LABEL: "WAN TX",
+        ATTR_GROUP: "Accounting",
+        ATTR_UNIT: "ps",
+        ATTR_UNIT_ATTR: "tx-rx-attr",
+        ATTR_PATH: "accounting",
+        ATTR_ATTR: "wan-tx",
+    },
+    "accounting_wan_rx": {
+        ATTR_DEVICE_CLASS: None,
+        ATTR_ICON: "mdi:upload-network",
+        ATTR_LABEL: "WAN RX",
+        ATTR_GROUP: "Accounting",
+        ATTR_UNIT: "ps",
+        ATTR_UNIT_ATTR: "tx-rx-attr",
+        ATTR_PATH: "accounting",
+        ATTR_ATTR: "wan-rx",
+    },
 }
 
+DEVICE_ATTRIBUTES_ACCOUNTING = [
+    "address",
+    "mac-address",
+    "comment"
+]
+
 
 # ---------------------------
 #   async_setup_entry
@@ -101,7 +160,7 @@ def update_items(inst, mikrotik_controller, async_add_entities, sensors):
     new_sensors = []
 
     for sensor in SENSOR_TYPES:
-        if "traffic_" not in sensor:
+        if "system_" in sensor:
             item_id = f"{inst}-{sensor}"
             _LOGGER.debug("Updating sensor %s", item_id)
             if item_id in sensors:
@@ -133,6 +192,24 @@ def update_items(inst, mikrotik_controller, async_add_entities, sensors):
                     )
                     new_sensors.append(sensors[item_id])
 
+        if "accounting_" in sensor:
+            for uid in mikrotik_controller.data["accounting"]:
+
+                item_id = f"{inst}-{sensor}-{mikrotik_controller.data['accounting'][uid]['mac-address']}"
+                if item_id in sensors:
+                    if sensors[item_id].enabled:
+                        sensors[item_id].async_schedule_update_ha_state()
+                    continue
+
+                if SENSOR_TYPES[sensor][ATTR_ATTR] in mikrotik_controller.data['accounting'][uid].keys():
+                    sensors[item_id] = MikrotikAccountingSensor(
+                        mikrotik_controller=mikrotik_controller,
+                        inst=inst,
+                        sensor=sensor,
+                        uid=uid,
+                    )
+                    new_sensors.append(sensors[item_id])
+
     if new_sensors:
         async_add_entities(new_sensors, True)
 
@@ -277,3 +354,65 @@ class MikrotikControllerTrafficSensor(MikrotikControllerSensor):
             self._data["default-name"],
             self._sensor,
         )
+
+
+# ---------------------------
+#   MikrotikAccountingSensor
+# ---------------------------
+class MikrotikAccountingSensor(MikrotikControllerSensor):
+    """Define an Mikrotik Accounting sensor."""
+
+    def __init__(self, mikrotik_controller, inst, sensor, uid):
+        """Initialize."""
+        super().__init__(mikrotik_controller, inst, sensor)
+        self._uid = uid
+        self._data = mikrotik_controller.data[SENSOR_TYPES[sensor][ATTR_PATH]][uid]
+
+    @property
+    def name(self):
+        """Return the name."""
+        return f"{self._inst} {self._data['host-name']} {self._type[ATTR_LABEL]} "
+
+    @property
+    def unique_id(self):
+        """Return a unique_id for this entity."""
+        return f"{self._inst.lower()}-{self._sensor.lower()}-{self._data['mac-address'].lower()}"
+
+    @property
+    def device_info(self):
+        """Return a accounting description for device registry."""
+        info = {
+            "identifiers": {
+                (
+                    DOMAIN,
+                    "serial-number",
+                    self._ctrl.data["routerboard"]["serial-number"],
+                    "sensor",
+                    "Accounting"
+                )
+            },
+            "manufacturer": self._ctrl.data["resource"]["platform"],
+            "model": self._ctrl.data["resource"]["board-name"],
+            "name": self._type[ATTR_GROUP],
+        }
+        return info
+
+    @property
+    def device_state_attributes(self):
+        """Return the state attributes."""
+        attributes = self._attrs
+        for variable in DEVICE_ATTRIBUTES_ACCOUNTING:
+            if variable in self._data:
+                attributes[format_attribute(variable)] = self._data[variable]
+
+        return attributes
+
+    async def async_added_to_hass(self):
+        """Port entity created."""
+        _LOGGER.debug(
+            "New sensor %s (%s [%s] %s)",
+            self._inst,
+            self._data["host-name"],
+            self._data["mac-address"],
+            self._sensor,
+        )