Make collector more resilient; validate records sequence nr.; enforce v9 version

Author: grafolean

dbutils.py

@@ -141,3 +141,7 @@ def migration_step_2():
 def migration_step_3():
     with db.cursor() as c:
         c.execute(f'CREATE INDEX {DB_PREFIX}records_ts ON {DB_PREFIX}records (ts);')
+
+def migration_step_4():
+    with db.cursor() as c:
+        c.execute(f'ALTER TABLE {DB_PREFIX}records DROP COLUMN version;')

docker-compose.yml

@@ -109,6 +109,14 @@ services:
       - netflowwriter
     volumes:
       - shared-grafolean:/shared-grafolean
+    # To use py-spy:
+    #   - $ docker exec -ti grafolean-newflow-writer bash
+    #   - # pip install py-spy
+    #   - # py-spy -n -o /tmp/prof/out.svg --pid 1
+    # But first, these 3 lines below must be enabled, to add a volume and capabilities: (careful not to add spaces!)
+    #  - /tmp/prof/:/tmp/prof/
+    #cap_add:
+    #  - SYS_PTRACE
     networks:
       - grafolean
 

netflowbot.py

@@ -16,7 +16,7 @@
 from dbutils import db, DB_PREFIX
 from lookup import PROTOCOLS
 
-logging.basicConfig(format='%(asctime)s | %(levelname)s | %(message)s',
+logging.basicConfig(format='%(asctime)s.%(msecs)03d | %(levelname)s | %(message)s',
                     datefmt='%Y-%m-%d %H:%M:%S', level=logging.DEBUG)
 logging.addLevelName(logging.DEBUG, color("DBG", 7))
 logging.addLevelName(logging.INFO, "INF")
@@ -90,10 +90,12 @@ def perform_job(*args, **job_params):
         output_path_prefix = f'entity.{entity_info["entity_id"]}.netflow'
 
         minute_ago = datetime.now() - timedelta(minutes=1)
+        two_minutes_ago = minute_ago - timedelta(minutes=1)
+
         values = []
         # Traffic in and out: (per interface)
-        values.extend(NetFlowBot.get_values_traffic_in(output_path_prefix, minute_ago))
-        values.extend(NetFlowBot.get_values_traffic_out(output_path_prefix, minute_ago))
+        values.extend(NetFlowBot.get_values_traffic_in(output_path_prefix, two_minutes_ago, minute_ago))
+        values.extend(NetFlowBot.get_values_traffic_out(output_path_prefix, two_minutes_ago, minute_ago))
 
         if not values:
             log.warning("No values found to be sent to Grafolean")
@@ -108,7 +110,7 @@ def perform_job(*args, **job_params):
         )
 
     @staticmethod
-    def get_values_traffic_in(output_path_prefix, from_time):
+    def get_values_traffic_in(output_path_prefix, from_time, to_time):
         with db.cursor() as c:
             # TODO: missing check for IP: r.client_ip = %s AND
             c.execute(f"""
@@ -120,11 +122,12 @@ def get_values_traffic_in(output_path_prefix, from_time):
                     {DB_PREFIX}flows "f"
                 WHERE
                     r.ts >= %s AND
+                    r.ts < %s AND
                     r.seq = f.record AND
                     ((f.data->'DIRECTION')::integer) = 1
                 GROUP BY
                     f.data->'INPUT_SNMP'
-            """, (from_time,))
+            """, (from_time, to_time,))
 
             values = []
             for interface_index, traffic_bytes in c.fetchall():
@@ -136,7 +139,7 @@ def get_values_traffic_in(output_path_prefix, from_time):
             return values
 
     @staticmethod
-    def get_values_traffic_out(output_path_prefix, from_time):
+    def get_values_traffic_out(output_path_prefix, from_time, to_time):
         with db.cursor() as c:
             # TODO: missing check for IP: r.client_ip = %s AND
             c.execute(f"""
@@ -148,11 +151,12 @@ def get_values_traffic_out(output_path_prefix, from_time):
                     {DB_PREFIX}flows "f"
                 WHERE
                     r.ts >= %s AND
+                    r.ts < %s AND
                     r.seq = f.record AND
                     ((f.data->'DIRECTION')::integer) = 1
                 GROUP BY
                     f.data->'OUTPUT_SNMP'
-            """, (from_time,))
+            """, (from_time, to_time,))
 
             values = []
             for interface_index, traffic_bytes in c.fetchall():

netflowcollector.py

@@ -15,7 +15,7 @@
 from pynetflow.main import get_export_packets
 
 
-logging.basicConfig(format='%(asctime)s | %(levelname)s | %(message)s',
+logging.basicConfig(format='%(asctime)s.%(msecs)03d | %(levelname)s | %(message)s',
                     datefmt='%Y-%m-%d %H:%M:%S', level=logging.DEBUG)
 logging.addLevelName(logging.DEBUG, color("DBG", 7))
 logging.addLevelName(logging.INFO, "INF")
@@ -26,16 +26,53 @@
 
 def process_netflow(netflow_port, named_pipe_filename):
     # endless loop - read netflow packets, encode them to JSON and write them to named pipe:
-    with open(named_pipe_filename, "wb", 0) as fp:
-        for ts, client, export in get_export_packets('0.0.0.0', NETFLOW_PORT):
-            entry = {
-                "ts": ts,
-                "client": client,
-                "version": export.header.version,
-                "flows": [flow.data for flow in export.flows],
-            }
-            line = json.dumps(entry).encode() + b'\n'
-            fp.write(line)
+    line = None
+    last_record_seqs = {}
+    while True:
+        try:
+            with open(named_pipe_filename, "wb", 0) as fp:
+                # if named pipe threq an error for some reason (BrokenPipe), write the line we
+                # have in buffer before listening to new packets:
+                if line is not None:
+                    fp.write(line)
+                    line = None
+                for ts, client, export in get_export_packets('0.0.0.0', NETFLOW_PORT):
+                    if export.header.version != 9:
+                        log.error(f"Only Netflow v9 currently supported, ignoring record (version: [{export.header.version}])")
+                        continue
+
+                    client_ip, _ = client
+
+                    # check for missing records:
+                    last_record_seq = last_record_seqs.get(client_ip)
+                    if last_record_seq is None:
+                        log.warning(f"Last record sequence number is not known, starting with {export.header.sequence}")
+                    elif export.header.sequence != last_record_seq + 1:
+                        log.error(f"Sequence number ({export.header.sequence}) does not follow ({last_record_seq}), some records might have been skipped")
+                    last_record_seqs[client_ip] = export.header.sequence
+
+                    flows_data = [flow.data for flow in export.flows]
+                    entry = {
+                        "ts": ts,
+                        "client": client_ip,
+                        "flows": [{
+                            "IN_BYTES": data["IN_BYTES"],
+                            "PROTOCOL": data["PROTOCOL"],
+                            "DIRECTION": data["DIRECTION"],
+                            "INPUT_SNMP": data["INPUT_SNMP"],
+                            "L4_DST_PORT": data["L4_DST_PORT"],
+                            "L4_SRC_PORT": data["L4_SRC_PORT"],
+                            "OUTPUT_SNMP": data["OUTPUT_SNMP"],
+                            "IPV4_DST_ADDR": data["IPV4_DST_ADDR"],
+                            "IPV4_SRC_ADDR": data["IPV4_SRC_ADDR"],
+                        } for data in flows_data],
+                    }
+                    line = json.dumps(entry).encode() + b'\n'
+                    fp.write(line)
+                    line = None
+        except Exception as ex:
+            log.exception(f"Exception: {str(ex)}")
+
 
 
 if __name__ == "__main__":

netflowwriter.py

@@ -16,7 +16,7 @@
 from dbutils import migrate_if_needed, db, DB_PREFIX
 
 
-logging.basicConfig(format='%(asctime)s | %(levelname)s | %(message)s',
+logging.basicConfig(format='%(asctime)s.%(msecs)03d | %(levelname)s | %(message)s',
                     datefmt='%Y-%m-%d %H:%M:%S', level=logging.DEBUG)
 logging.addLevelName(logging.DEBUG, color("DBG", 7))
 logging.addLevelName(logging.INFO, "INF")
@@ -46,9 +46,9 @@ def process_named_pipe(named_pipe_filename):
 def write_record(j):
     with db.cursor() as c:
         # first save the flow record:
-        ts = datetime.utcfromtimestamp(j["ts"])
-        log.info(f"Received record: {ts} from {j['client'][0]}")
-        c.execute(f"INSERT INTO {DB_PREFIX}records (ts, client_ip, version) VALUES (%s, %s, %s) RETURNING seq;", (ts, j['client'][0], j['version'],))
+        ts = datetime.utcfromtimestamp(j['ts'])
+        log.info(f"Received record: {ts} from {j['client']}")
+        c.execute(f"INSERT INTO {DB_PREFIX}records (ts, client_ip) VALUES (%s, %s) RETURNING seq;", (ts, j['client'],))
         record_seq = c.fetchone()[0]
 
         # then save each of the flows within the record, but use execute_values() to perform bulk insert: