fixed IPv6 connection parsing / added separate IPv4/IPv6 connection totals

Author: akpw

mktxp/collector/connection_collector.py

@@ -27,6 +27,10 @@ def collect(router_entry):
             if connection_records:
                 connection_metrics = BaseCollector.gauge_collector('ip_connections_total', 'Number of IP connections', connection_records, 'count',)
                 yield connection_metrics
+                ipv4_connection_metrics = BaseCollector.gauge_collector('ipv4_connections_total', 'Number of IPv4 connections', connection_records, 'ipv4_count',)
+                yield ipv4_connection_metrics
+                ipv6_connection_metrics = BaseCollector.gauge_collector('ipv6_connections_total', 'Number of IPv6 connections', connection_records, 'ipv6_count',)
+                yield ipv6_connection_metrics
 
         if router_entry.config_entry.connection_stats:
             connection_stats_records = IPConnectionStatsDatasource.metric_records(router_entry)

mktxp/datasource/connection_ds.py

@@ -24,14 +24,38 @@ def metric_records(router_entry, *, metric_labels = None):
         if metric_labels is None:
             metric_labels = []
         try:
-            res = router_entry.api_connection.router_api().get_resource('/ip/firewall/connection/').call('print', {'count-only': ''})
-            # result processing as described at: https://github.com/socialwifi/RouterOS-api/issues/79#issuecomment-2089744809
+            api = router_entry.api_connection.router_api()
+            ipv4_cnt_str = '0'
+            ipv6_cnt_str = '0'
+
+            res = api.get_resource('/ip/firewall/connection/').call('print', {'count-only': ''})
             cnt_str = res.done_message.get('ret')
+            if cnt_str is not None:
+                ipv4_cnt_str = cnt_str
+
+            try:
+                res_v6 = api.get_resource('/ipv6/firewall/connection/').call('print', {'count-only': ''})
+                cnt_str_v6 = res_v6.done_message.get('ret')
+                if cnt_str_v6 is not None:
+                    ipv6_cnt_str = cnt_str_v6
+            except Exception:
+                pass
+
+            try:
+                ipv4_count = int(ipv4_cnt_str)
+            except (ValueError, TypeError):
+                ipv4_count = 0
+
             try:
-                count = int(cnt_str)
+                ipv6_count = int(ipv6_cnt_str)
             except (ValueError, TypeError):
-                cnt_str = '0'
-            records = [{'count': cnt_str}]
+                ipv6_count = 0
+
+            records = [{
+                'count': str(ipv4_count + ipv6_count),
+                'ipv4_count': str(ipv4_count),
+                'ipv6_count': str(ipv6_count)
+            }]
             return BaseDSProcessor.trimmed_records(router_entry, router_records = records, metric_labels = metric_labels)
         except Exception as exc:
             print(f'Error getting IP connection info from router {router_entry.router_name}@{router_entry.config_entry.hostname}: {exc}')
@@ -41,6 +65,15 @@ def metric_records(router_entry, *, metric_labels = None):
 class IPConnectionStatsDatasource:
     ''' IP connections stats data provider
     '''
+    @staticmethod
+    def strip_port(address):
+        if address.startswith('['):
+            return address[1:address.find(']')]
+        colons = address.count(':')
+        if colons == 1:
+            return address.split(':')[0]
+        return address
+
     @staticmethod
     def metric_records(router_entry, *, metric_labels = None, add_router_id = True):
         if metric_labels is None:
@@ -56,12 +89,18 @@ def metric_records(router_entry, *, metric_labels = None, add_router_id = True):
             else:
                 proplist = 'src-address'
 
-            connection_records = router_entry.api_connection.router_api().get_resource('/ip/firewall/connection/').call('print', \
-                                                                                            {'proplist': proplist})
+            api = router_entry.api_connection.router_api()
+            connection_records = api.get_resource('/ip/firewall/connection/').call('print', {'proplist': proplist})
+            try:
+                connection_records_v6 = api.get_resource('/ipv6/firewall/connection/').call('print', {'proplist': proplist})
+                connection_records.extend(connection_records_v6)
+            except Exception:
+                pass
+
              # calculate number of connections per src-address
             connections_per_src_address = {}
             for connection_record in connection_records:
-                address = connection_record['src-address'].split(':')[0]
+                address = IPConnectionStatsDatasource.strip_port(connection_record.get('src-address', ''))
 
                 count, destinations = 0, set()
                 if connections_per_src_address.get(address):

tests/collector/test_connection_collector.py

@@ -55,14 +55,14 @@ def api_call_router(*args, **kwargs):
     # Test the method of focus
     result = IPConnectionStatsDatasource.metric_records(mock_router_entry)
     if should_make_stats_call:
-        # This one should have been called twice, once for count and once for the stats
-        assert call_mock.call_count == 2
+        # This one should have been called four times, twice for count (IPv4+IPv6) and twice for the stats (IPv4+IPv6)
+        assert call_mock.call_count == 4
         assert result is not None
         assert len(result) > 0
         assert result[0]['src_address'] == '1.1.1.1'
     else:
-        # And this just once for the count
-        call_mock.assert_called_once()
+        # And this twice for the count (IPv4+IPv6)
+        assert call_mock.call_count == 2
         assert result == []
 
 def test_ip_connection_stats_datasource_without_destinations():
@@ -100,9 +100,19 @@ def api_call_router(*args, **kwargs):
 
     result = IPConnectionStatsDatasource.metric_records(mock_router_entry)
 
-    assert call_mock.call_count == 2
+    assert call_mock.call_count == 4
     assert result is not None
     assert len(result) == 1
     assert result[0]['src_address'] == '1.1.1.1'
-    assert result[0]['connection_count'] == 2
+    assert result[0]['connection_count'] == 4
     assert result[0]['dst_addresses'] == ''
+
+def test_strip_port_ipv4_and_ipv6():
+    """
+    Verifies that IPConnectionStatsDatasource.strip_port correctly parses IPv4 and IPv6 addresses.
+    """
+    assert IPConnectionStatsDatasource.strip_port('192.168.1.1:80') == '192.168.1.1'
+    assert IPConnectionStatsDatasource.strip_port('192.168.1.1') == '192.168.1.1'
+    assert IPConnectionStatsDatasource.strip_port('[2001:db8::1]:443') == '2001:db8::1'
+    assert IPConnectionStatsDatasource.strip_port('2001:db8::1') == '2001:db8::1'
+    assert IPConnectionStatsDatasource.strip_port('') == ''