Merge pull request #2 from grafolean/fix/ipv6

Fix/ipv6

Author: grafolean

dbutils.py

@@ -158,4 +158,10 @@ def migration_step_3():
 def migration_step_4():
     with get_db_cursor() as c:
         # INTEGER (signed 32 bits) is not enough for in_bytes:
-        c.execute(f"ALTER TABLE {DB_PREFIX}flows ALTER COLUMN in_bytes SET DATA TYPE BIGINT;")
+        c.execute(f"ALTER TABLE {DB_PREFIX}flows ALTER COLUMN in_bytes SET DATA TYPE BIGINT;")
+
+def migration_step_5():
+    """ We want to save IPv6 addresses too """
+    with get_db_cursor() as c:
+        c.execute(f"ALTER TABLE {DB_PREFIX}flows RENAME COLUMN ipv4_dst_addr TO ipvX_dst_addr;")
+        c.execute(f"ALTER TABLE {DB_PREFIX}flows RENAME COLUMN ipv4_src_addr TO ipvX_src_addr;")

netflowbot.py

@@ -50,7 +50,7 @@ def wrapper(*args, **kwargs):
 
 
 def path_part_encode(s):
-    return s.replace(".", '%2e')
+    return s.replace(".", '%2e').replace(":", '%3a')
 
 
 def _get_last_used_ts(job_id):
@@ -195,8 +195,8 @@ def perform_account_aggr_job(*args, **job_params):
         #   l4_src_port   | integer       | source port
         #   input_snmp    | smallint      | input interface index
         #   output_snmp   | smallint      | output interface index
-        #   ipv4_dst_addr | inet          | source IP
-        #   ipv4_src_addr | inet          | destination IP
+        #   ipvX_dst_addr | inet          | source IP
+        #   ipvX_src_addr | inet          | destination IP
 
         job_id = job_params["job_id"]
         interval_label = job_params["interval_label"]
@@ -348,7 +348,7 @@ def get_top_N_IPs_for_entity_interfaces(interval_label, last_used_ts, max_ts, ti
             for interface_index, in c.fetchall():
                 c2.execute(f"""
                     SELECT
-                        f.{'ipv4_src_addr' if direction == DIRECTION_INGRESS else 'ipv4_dst_addr'},
+                        f.{'ipvX_src_addr' if direction == DIRECTION_INGRESS else 'ipvX_dst_addr'},
                         sum(f.in_bytes) "traffic"
                     FROM
                         {DB_PREFIX}flows "f"
@@ -359,7 +359,7 @@ def get_top_N_IPs_for_entity_interfaces(interval_label, last_used_ts, max_ts, ti
                         f.direction = %s AND
                         f.{'input_snmp' if direction == DIRECTION_INGRESS else 'output_snmp'} = %s
                     GROUP BY
-                        f.{'ipv4_src_addr' if direction == DIRECTION_INGRESS else 'ipv4_dst_addr'}
+                        f.{'ipvX_src_addr' if direction == DIRECTION_INGRESS else 'ipvX_dst_addr'}
                     ORDER BY
                         traffic desc
                     LIMIT {TOP_N_MAX};
@@ -382,7 +382,7 @@ def get_top_N_IPs_for_entity(interval_label, last_used_ts, max_ts, time_between,
             values = []
             c.execute(f"""
                 SELECT
-                    f.{'ipv4_src_addr' if direction == DIRECTION_INGRESS else 'ipv4_dst_addr'},
+                    f.{'ipvX_src_addr' if direction == DIRECTION_INGRESS else 'ipvX_dst_addr'},
                     sum(f.in_bytes) "traffic"
                 FROM
                     {DB_PREFIX}flows "f"
@@ -392,7 +392,7 @@ def get_top_N_IPs_for_entity(interval_label, last_used_ts, max_ts, time_between,
                     f.ts <= %s AND
                     f.direction = %s
                 GROUP BY
-                    f.{'ipv4_src_addr' if direction == DIRECTION_INGRESS else 'ipv4_dst_addr'}
+                    f.{'ipvX_src_addr' if direction == DIRECTION_INGRESS else 'ipvX_dst_addr'}
                 ORDER BY
                     traffic desc
                 LIMIT {TOP_N_MAX};
@@ -499,7 +499,7 @@ def get_top_N_connections_for_entity(interval_label, last_used_ts, max_ts, time_
             values = []
             c.execute(f"""
                 SELECT
-                    f.ipv4_src_addr, f.ipv4_dst_addr,
+                    f.ipvX_src_addr, f.ipvX_dst_addr,
                     sum(f.in_bytes) "traffic"
                 FROM
                     {DB_PREFIX}flows "f"
@@ -509,15 +509,15 @@ def get_top_N_connections_for_entity(interval_label, last_used_ts, max_ts, time_
                     f.ts <= %s AND
                     f.direction = %s
                 GROUP BY
-                    f.ipv4_src_addr, f.ipv4_dst_addr
+                    f.ipvX_src_addr, f.ipvX_dst_addr
                 ORDER BY
                     traffic desc
                 LIMIT {TOP_N_MAX};
             """, (entity_ip, last_used_ts, max_ts, direction,))
 
             output_path_entity = NetFlowBot.construct_output_path_prefix(interval_label, direction, entity_id, interface=None)
-            for ipv4_src_addr, ipv4_dst_addr, traffic_bytes in c.fetchall():
-                output_path = f"{output_path_entity}.topconn.{path_part_encode(ipv4_src_addr)}.{path_part_encode(ipv4_dst_addr)}"
+            for ipvX_src_addr, ipvX_dst_addr, traffic_bytes in c.fetchall():
+                output_path = f"{output_path_entity}.topconn.{path_part_encode(ipvX_src_addr)}.{path_part_encode(ipvX_dst_addr)}"
                 values.append({
                     'p': output_path,
                     'v': traffic_bytes / time_between,  # Bps

netflowwriter.py

@@ -41,9 +41,13 @@
 # 11-byte signature (constructed in this way to detect possible mangled bytes), flags, header extension
 # https://www.postgresql.org/docs/9.0/sql-copy.html#AEN59377
 PG_COPYFROM_INIT = struct.pack('!11sII', b'PGCOPY\n\377\r\n\0', 0, 0)
-# 4-byte INETv4 prefix: family, netmask, is_cidr, n bytes
+# 4-byte INETv4/v6 prefix: family, netmask, is_cidr, n bytes
 # https://doxygen.postgresql.org/network_8c_source.html#l00193
-IPV4_PREFIX = struct.pack('!BBBB', socket.AF_INET, 32, 0, 4)
+IPV4_ADDRESS_PREFIX = struct.pack('!BBBB', socket.AF_INET, 32, 0, 4)
+# Gotcha: IPv6 address family in Postgres is *not* socket.AF_INET6 (10),
+# instead it is defined as socket.AF_INET + 1 (2 + 1 == 3)
+# https://doxygen.postgresql.org/utils_2inet_8h_source.html#l00040
+IPV6_ADDRESS_PREFIX = struct.pack('!BBBB', socket.AF_INET + 1, 128, 0, 16)
 
 
 def _pgwriter_init():
@@ -52,22 +56,30 @@ def _pgwriter_init():
     return pg_writer
 
 
-def _pgwriter_write(pgwriter, ts, client_ip, IN_BYTES, PROTOCOL, DIRECTION, L4_DST_PORT, L4_SRC_PORT, INPUT_SNMP, OUTPUT_SNMP, IPV4_DST_ADDR, IPV4_SRC_ADDR):
-    buf = struct.pack('!HiIi4s4siQiHiHiIiIiHiHi4s4si4s4s',
+def _pgwriter_write(pgwriter, ts, client_ip, IN_BYTES, PROTOCOL, DIRECTION, L4_DST_PORT, L4_SRC_PORT, INPUT_SNMP, OUTPUT_SNMP, address_family, IPVx_DST_ADDR, IPVx_SRC_ADDR):
+    buf = struct.pack('!HiIi4s4siQiHiHiIiIiHiH',
         11,  # number of columns
         4, int(ts),                       # integer - beware of Y2038 problem! :)
-        8, IPV4_PREFIX, socket.inet_aton(client_ip),   # 4 bytes prefix + 4 bytes IP
+        8, IPV4_ADDRESS_PREFIX, socket.inet_aton(client_ip),   # 4 bytes prefix + 4 bytes IP
         8, IN_BYTES,                      # bigint
         2, PROTOCOL,
         2, DIRECTION,
         4, L4_DST_PORT,
         4, L4_SRC_PORT,
         2, INPUT_SNMP,
         2, OUTPUT_SNMP,
-        8, IPV4_PREFIX, IPV4_DST_ADDR,
-        8, IPV4_PREFIX, IPV4_SRC_ADDR,
     )
-    pgwriter.write(buf)
+    if address_family != socket.AF_INET6:
+        buf2 = struct.pack('!i4s4si4s4s',
+            8, IPV4_ADDRESS_PREFIX, IPVx_DST_ADDR,
+            8, IPV4_ADDRESS_PREFIX, IPVx_SRC_ADDR,
+        )
+    else:
+        buf2 = struct.pack('!i4s16si4s16s',
+            4 + 16, IPV6_ADDRESS_PREFIX, IPVx_DST_ADDR,
+            4 + 16, IPV6_ADDRESS_PREFIX, IPVx_SRC_ADDR,
+        )
+    pgwriter.write(buf + buf2)
 
 
 def _pgwriter_finish(pgwriter):
@@ -138,9 +150,8 @@ def process_named_pipe(named_pipe_filename):
                     except UnknownNetFlowVersion:
                         log.warning("Unknown NetFlow version")
                         continue
-                    except TemplateNotRecognized:
-                        log.warning("Failed to decode a v9 ExportPacket, template not "
-                            "recognized (if this happens at the start, it's ok)")
+                    except TemplateNotRecognized as ex:
+                        log.warning(f"Failed to decode a v9 ExportPacket, template not recognized (if this happens at the start, it's ok). Template id: {ex.template_id}")
                         continue
 
                 except Exception as ex:
@@ -194,39 +205,36 @@ def write_buffer(buffer, partition_no):
 
 
     log.debug(f"Writing {len(buffer)} records to DB, partition {partition_no}")
-    ipv6_ignored_records = 0  # we don't support IPv6 yet
     # save each of the flows within the record, but use execute_values() to perform bulk insert:
     def _get_data(buffer):
         for ts, client_ip, export in buffer:
             netflow_version, flows = export.header.version, export.flows
             if netflow_version == 9:
                 for f in flows:
                     try:
-                        if f.data.get("IP_PROTOCOL_VERSION", 4) == 6:
-                            ipv6_ignored_records += 1
-                            continue
+                        # if f.data.get("IP_PROTOCOL_VERSION", 4) == 6:
+                        if not f.data.get("IPV6_DST_ADDR", None) is None:
+                            address_family = socket.AF_INET6
+                            dst = socket.inet_pton(address_family, f.data["IPV6_DST_ADDR"])
+                            src = socket.inet_pton(address_family, f.data["IPV6_SRC_ADDR"])
+                        else:
+                            address_family = socket.AF_INET
+                            dst = socket.inet_aton(f.data["IPV4_DST_ADDR"])
+                            src = socket.inet_aton(f.data["IPV4_SRC_ADDR"])
 
                         yield (
                             ts,
                             client_ip,
-                            # "IN_BYTES":
                             f.data["IN_BYTES"],
-                            # "PROTOCOL":
                             f.data["PROTOCOL"],
-                            # "DIRECTION":
                             f.data.get("DIRECTION", DIRECTION_INGRESS),
-                            # "L4_DST_PORT":
                             f.data["L4_DST_PORT"],
-                            # "L4_SRC_PORT":
                             f.data["L4_SRC_PORT"],
-                            # "INPUT_SNMP":
                             f.data["INPUT_SNMP"],
-                            # "OUTPUT_SNMP":
                             f.data["OUTPUT_SNMP"],
-                            # "IPV4_DST_ADDR":
-                            socket.inet_aton(f.data["IPV4_DST_ADDR"]),
-                            # "IPV4_SRC_ADDR":
-                            socket.inet_aton(f.data["IPV4_SRC_ADDR"]),
+                            address_family,
+                            dst,
+                            src,
                         )
                     except KeyError:
                         log.exception(f"[{client_ip}] Error decoding v9 flow. Contents: {repr(f.data)}")
@@ -250,6 +258,8 @@ def _get_data(buffer):
                             f.data["INPUT"],
                             # "OUTPUT_SNMP":
                             f.data["OUTPUT"],
+                            # address_family is always IPv4:
+                            socket.AF_INET,
                             # netflow v5 IP addresses are decoded to integers, which is less suitable for us - pack
                             # them back to bytes:
                             # "IPV4_DST_ADDR":
@@ -267,9 +277,6 @@ def _get_data(buffer):
         _pgwriter_write(pgwriter, *data)
     _pgwriter_finish(pgwriter)
 
-    if ipv6_ignored_records > 0:
-        log.error(f"We do not support IPv6 (yet), some IPv6 flow records were ignored: {ipv6_ignored_records}")
-
 
 if __name__ == "__main__":
     NAMED_PIPE_FILENAME = os.environ.get('NAMED_PIPE_FILENAME', None)