Merge pull request #23 from ivanpavlina/master

Adding accounting monitoring

README.md

@@ -25,7 +25,8 @@ Features:
  * Configurable update interval
  * Configurable traffic unit (bps, Kbps, Mbps, B/s, KB/s, MB/s)
  * Supports monitoring of multiple mikrotik devices simultaneously
-
+ * RX/TX WAN/LAN traffic sensors per hosts from Mikrotik Accounting feature
+ 
 # Integration preview
 ![Tracker and sensors](https://raw.githubusercontent.com/tomaae/homeassistant-mikrotik_router/master/docs/assets/images/ui/device_tracker.png)
 
@@ -39,6 +40,7 @@ Features:
 ![Queue switch](https://raw.githubusercontent.com/tomaae/homeassistant-mikrotik_router/master/docs/assets/images/ui/queue_switch.png)
 
 ![Host tracker](https://raw.githubusercontent.com/tomaae/homeassistant-mikrotik_router/master/docs/assets/images/ui/host_tracker.png)
+![Accounting sensor](https://raw.githubusercontent.com/tomaae/homeassistant-mikrotik_router/master/docs/assets/images/ui/accounting_sensor.png)
 
 # Setup integration
 Setup this integration for your Mikrotik device in Home Assistant via `Configuration -> Integrations -> Add -> Mikrotik Router`.
@@ -47,7 +49,7 @@ You can add this integration several times for different devices.
 ![Add Integration](https://raw.githubusercontent.com/tomaae/homeassistant-mikrotik_router/master/docs/assets/images/ui/setup_integration.png)
 * "Host" - Use hostname or IP
 * "Port" - Leave at 0 for defaults
-* "Name of the integration" - Friendy name for this router
+* "Name of the integration" - Friendly name for this router
 * "Unit of measurement" - Traffic sensor measurement (bps, Kbps, Mbps, B/s, KB/s, MB/s)
 
 # Configuration
@@ -58,3 +60,12 @@ You can add this integration several times for different devices.
 
 ## List of detected devices
 ![Integration options](https://raw.githubusercontent.com/tomaae/homeassistant-mikrotik_router/master/docs/assets/images/ui/integration_devices.png)
+
+## Accounting
+For per-IP throughput tracking Mikrotik's accounting feature is used.
+
+[Mikrotik support page](https://wiki.mikrotik.com/wiki/Manual:IP/Accounting)
+
+Feature will be automatically used if accounting is enabled in Mikrotik. Feature is present in Winbox IP-Accounting. Make sure that threshold is set to reasonable value to store all connections between user defined scan interval. Max value is 8192 so for piece of mind I recommend setting that value. Web Access is not needed, integration is using API access. 
+
+Integration will scan DHCP Lease table and ARP table to generate all known hosts. For every host aleast two sensors for WAN traffic (mikrotik-XXX-wan-rx and mikrotik-XXX-wan-tx) are created. If the parameter *account-local-traffic* is set in Mikrotik's accounting configuration it will also create two sensors for LAN traffic (mikrotik-XXX-lan-rx and mikrotik-XXX-lan-tx).

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/device_tracker.py

@@ -38,7 +38,7 @@ DEVICE_ATTRIBUTES_IFACE = [
 ]
 
 DEVICE_ATTRIBUTES_HOST = [
-    "hostname",
+    "host-name",
     "address",
     "mac-address",
     "interface",
@@ -234,7 +234,7 @@ class MikrotikControllerHostDeviceTracker(ScannerEntity):
         _LOGGER.debug(
             "New host tracker %s (%s - %s)",
             self._inst,
-            self._data["hostname"],
+            self._data["host-name"],
             self._data["mac-address"],
         )
 
@@ -254,7 +254,7 @@ class MikrotikControllerHostDeviceTracker(ScannerEntity):
     @property
     def name(self):
         """Return the name of the host."""
-        return f"{self._data['hostname']}"
+        return f"{self._data['host-name']}"
 
     @property
     def unique_id(self):

custom_components/mikrotik_router/mikrotik_controller.py

@@ -66,6 +66,7 @@ class MikrotikControllerData:
             "dhcp-network": {},
             "dhcp": {},
             "host": {},
+            "accounting": {}
         }
 
         self.listeners = []
@@ -205,6 +206,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)
 
@@ -561,7 +564,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"""
@@ -794,30 +797,30 @@ class MikrotikControllerData:
         for uid, vals in self.data["host"].items():
             # Add missing default values
             for key, default in zip(
-                    ["address", "mac-address", "interface", "hostname", "last-seen", "available"],
+                    ["address", "mac-address", "interface", "host-name", "last-seen", "available"],
                     ["unknown", "unknown", "unknown", "unknown", False],
             ):
                 if key not in self.data["host"][uid]:
                     self.data["host"][uid][key] = default
 
             # Resolve hostname
-            if vals["hostname"] == "unknown":
+            if vals["host-name"] == "unknown":
                 if vals["address"] != "unknown":
                     for dns_uid, dns_vals in self.data["dns"].items():
                         if dns_vals["address"] == vals["address"]:
-                            self.data["host"][uid]["hostname"] = dns_vals["name"].split('.')[0]
+                            self.data["host"][uid]["host-name"] = dns_vals["name"].split('.')[0]
                             break
 
-                if self.data["host"][uid]["hostname"] == "unknown" \
+                if self.data["host"][uid]["host-name"] == "unknown" \
                         and uid in self.data["dhcp"] and self.data["dhcp"][uid]["comment"] != "":
-                    self.data["host"][uid]["hostname"] = self.data["dhcp"][uid]["comment"]
+                    self.data["host"][uid]["host-name"] = self.data["dhcp"][uid]["comment"]
 
-                elif self.data["host"][uid]["hostname"] == "unknown" \
+                elif self.data["host"][uid]["host-name"] == "unknown" \
                         and uid in self.data["dhcp"] and self.data["dhcp"][uid]["host-name"] != "unknown":
-                    self.data["host"][uid]["hostname"] = self.data["dhcp"][uid]["host-name"]
+                    self.data["host"][uid]["host-name"] = self.data["dhcp"][uid]["host-name"]
 
-                elif self.data["host"][uid]["hostname"] == "unknown":
-                    self.data["host"][uid]["hostname"] = uid
+                elif self.data["host"][uid]["host-name"] == "unknown":
+                    self.data["host"][uid]["host-name"] = uid
 
             # Check host availability
             if vals["address"] != "unknown" and vals["interface"] != "unknown":
@@ -827,3 +830,118 @@ class MikrotikControllerData:
             # Update last seen
             if self.data["host"][uid]["available"]:
                 self.data["host"][uid]["last-seen"] = utcnow()
+
+    def _address_part_of_local_network(self, address):
+        address = ip_address(address)
+        for vals in self.data["dhcp-network"].values():
+            if address in vals["IPv4Network"]:
+                return True
+        return False
+
+    def _get_accounting_uid_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()
+
+        # Build missing hosts from main hosts dict
+        for uid, vals in self.data["host"].items():
+            if uid not in self.data["accounting"]:
+                self.data["accounting"][uid] = {
+                    'address': vals['address'],
+                    'mac-address': vals['mac-address'],
+                    'host-name': vals['host-name'],
+                    'tx-rx-attr': traffic_type,
+                    'available': False,
+                    'local_accounting': False
+                }
+
+        _LOGGER.debug(f"Working with {len(self.data['accounting'])} accounting devices")
+
+        # Build temp accounting values dict with ip address as key
+        tmp_accounting_values = {}
+        for uid, vals in self.data['accounting'].items():
+            tmp_accounting_values[vals['address']] = {
+                "wan-tx": 0,
+                "wan-rx": 0,
+                "lan-tx": 0,
+                "lan-rx": 0
+            }
+
+        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
+
+        # Calculate real throughput and transform it to appropriate unit
+        # Also handle availability of accounting and local_accounting from Mikrotik
+        for addr in tmp_accounting_values:
+            uid = self._get_accounting_uid_by_ip(addr)
+            if not uid:
+                _LOGGER.warning(f"Address {addr} not found in accounting data, skipping update")
+                continue
+
+            self.data['accounting'][uid]['tx-rx-attr'] = traffic_type
+            self.data['accounting'][uid]['available'] = accounting_enabled
+            self.data['accounting'][uid]['local_accounting'] = local_traffic_enabled
+
+            if not accounting_enabled:
+                # Skip calculation for WAN and LAN, accounting is disabled
+                continue
+
+            self.data['accounting'][uid]['wan-tx'] = round(
+                tmp_accounting_values[addr]['wan-tx'] / time_diff * traffic_div, 2) \
+                if tmp_accounting_values[addr]['wan-tx'] else 0.0
+
+            self.data['accounting'][uid]['wan-rx'] = round(
+                tmp_accounting_values[addr]['wan-rx'] / time_diff * traffic_div, 2) \
+                if tmp_accounting_values[addr]['wan-rx'] else 0.0
+
+            if not local_traffic_enabled:
+                # Skip calculation for LAN, LAN accounting is disabled
+                continue
+
+            self.data['accounting'][uid]['lan-tx'] = round(
+                tmp_accounting_values[addr]['lan-tx'] / time_diff * traffic_div, 2) \
+                if tmp_accounting_values[addr]['lan-tx'] else 0.0
+
+            self.data['accounting'][uid]['lan-rx'] = round(
+                tmp_accounting_values[addr]['lan-rx'] / time_diff * traffic_div, 2) \
+                if tmp_accounting_values[addr]['lan-rx'] else 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:upload-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:download-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:upload-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:download-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",
+    "host-name"
+]
+
 
 # ---------------------------
 #   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,75 @@ 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 available(self) -> bool:
+        """Return if controller and accounting feature in Mikrotik is available.
+           Additional check for lan-tx/rx sensors
+        """
+        if self._attr in ['lan-tx', 'lan-rx']:
+            return self._ctrl.connected() and self._data['available'] and self._data['local_accounting']
+        else:
+            return self._ctrl.connected() and self._data['available']
+
+    @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,
+        )