Merge branch 'props-master'

Merging pull request #9 by @pR0Ps
https://github.com/bitkeks/python-netflow-v9-softflowd/pull/9

Thanks for the contribution!

Resolves #9

README.md

@@ -1,5 +1,5 @@
 # Python NetFlow v9 parser and UDP collector
-This script is able to parse incoming UDP NetFlow packets of **NetFlow version 9**.
+This script is able to collect and parse incoming UDP NetFlow packets of **NetFlow versions 1, 5 and 9**.
 
 Version 9 is the first NetFlow version using templates.
 Templates make dynamically sized and configured NetFlow data flowsets possible,

analyze_json.py

@@ -8,33 +8,30 @@ Copyright 2017-2019 Dominik Pataky <dev@bitkeks.eu>
 Licensed under MIT License. See LICENSE.
 """
 
+import argparse
+from collections import namedtuple
+import contextlib
 from datetime import datetime
+import functools
 import ipaddress
 import json
-import os.path
-import sys
 import socket
-from collections import namedtuple
+import sys
 
-Pair = namedtuple('Pair', 'src dest')
 
-def getIPs(flow):
+Pair = namedtuple('Pair', ['src', 'dest'])
-    use_ipv4 = False  # optimistic default case of IPv6
 
-    if 'IP_PROTOCOL_VERSION' in flow and flow['IP_PROTOCOL_VERSION'] == 4:
-        use_ipv4 = True
-    elif 'IPV4_SRC_ADDR' in flow or 'IPV4_DST_ADDR' in flow:
-        use_ipv4 = True
 
-    if use_ipv4:
+@functools.lru_cache(maxsize=128)
-        return Pair(
+def resolve_hostname(ip):
-            ipaddress.ip_address(flow['IPV4_SRC_ADDR']),
+    return socket.getfqdn(ip)
-            ipaddress.ip_address(flow['IPV4_DST_ADDR']))
 
-    # else: return IPv6 pair
+
-    return Pair(
+def fallback(d, keys):
-        ipaddress.ip_address(flow['IPV6_SRC_ADDR']),
+    for k in keys:
-        ipaddress.ip_address(flow['IPV6_DST_ADDR']))
+        if k in d:
+            return d[k]
+    raise KeyError(", ".join(keys))
 
 
 class Connection:
@@ -42,33 +39,62 @@ class Connection:
     The direction of the data flow can be seen by looking at the size.
 
     'src' describes the peer which sends more data towards the other. This
-    does NOT have to mean, that 'src' was the initiator of the connection.
+    does NOT have to mean that 'src' was the initiator of the connection.
     """
     def __init__(self, flow1, flow2):
-        if flow1['IN_BYTES'] >= flow2['IN_BYTES']:
+        if not flow1 or not flow2:
+            raise Exception("A connection requires two flows")
+
+        # Assume the size that sent the most data is the source
+        # TODO: this might not always be right, maybe use earlier timestamp?
+        size1 = fallback(flow1, ['IN_BYTES', 'IN_OCTETS'])
+        size2 = fallback(flow2, ['IN_BYTES', 'IN_OCTETS'])
+        if size1 >= size2:
             src = flow1
             dest = flow2
         else:
             src = flow2
             dest = flow1
 
-        ips = getIPs(src)
+        ips = self.get_ips(src)
         self.src = ips.src
         self.dest = ips.dest
-        self.src_port = src['L4_SRC_PORT']
+        self.src_port = fallback(src, ['L4_SRC_PORT', 'SRC_PORT'])
-        self.dest_port = src['L4_DST_PORT']
+        self.dest_port = fallback(dest, ['L4_DST_PORT', 'DST_PORT'])
-        self.size = src['IN_BYTES']
+        self.size = fallback(src, ['IN_BYTES', 'IN_OCTETS'])
 
         # Duration is given in milliseconds
         self.duration = src['LAST_SWITCHED'] - src['FIRST_SWITCHED']
         if self.duration < 0:
             # 32 bit int has its limits. Handling overflow here
+            # TODO: Should be handled in the collection phase
             self.duration = (2**32 - src['FIRST_SWITCHED']) + src['LAST_SWITCHED']
 
     def __repr__(self):
         return "<Connection from {} to {}, size {}>".format(
             self.src, self.dest, self.human_size)
 
+    @staticmethod
+    def get_ips(flow):
+        # TODO: These values should be parsed into strings in the collection phase.
+        #       The floating point representation of an IPv6 address in JSON
+        #       could lose precision.
+
+        # IPv4
+        if flow.get('IP_PROTOCOL_VERSION') == 4 \
+                or 'IPV4_SRC_ADDR' in flow \
+                or 'IPV4_DST_ADDR' in flow:
+            return Pair(
+                ipaddress.ip_address(flow['IPV4_SRC_ADDR']),
+                ipaddress.ip_address(flow['IPV4_DST_ADDR'])
+            )
+
+        # IPv6
+        return Pair(
+            ipaddress.ip_address(flow['IPV6_SRC_ADDR']),
+            ipaddress.ip_address(flow['IPV6_DST_ADDR'])
+        )
+
     @property
     def human_size(self):
         # Calculate a human readable size of the traffic
@@ -96,52 +122,42 @@ class Connection:
     @property
     def hostnames(self):
         # Resolve the IPs of this flows to their hostname
-        src_hostname = socket.getfqdn(self.src.compressed)
+        src_hostname = resolve_hostname(self.src.compressed)
-        dest_hostname = socket.getfqdn(self.dest.compressed)
+        dest_hostname = resolve_hostname(self.dest.compressed)
-
         return Pair(src_hostname, dest_hostname)
 
     @property
     def service(self):
         # Resolve ports to their services, if known
-        service = "unknown"
+        # Try source port, fallback to dest port, otherwise "unknown"
-        try:
+        with contextlib.suppress(OSError):
-            # Try service of sending peer first
+            return socket.getservbyport(self.src_port)
-            service = socket.getservbyport(self.src_port)
+        with contextlib.suppress(OSError):
-        except OSError:
+            return socket.getservbyport(self.dest_port)
-            # Resolving the sport did not work, trying dport
+        return "unknown"
-            try:
-                service = socket.getservbyport(self.dest_port)
-            except OSError:
-                pass
-        return service
 
 
-# Handle CLI args and load the data dump
+if __name__ == "__main__":
-if len(sys.argv) < 2:
+    parser = argparse.ArgumentParser(description="Output a basic analysis of NetFlow data")
-    exit("Use {} <filename>.json".format(sys.argv[0]))
+    parser.add_argument('filename', nargs='?', type=argparse.FileType('r'),
-filename = sys.argv[1]
+                        default=sys.stdin,
-if not os.path.exists(filename):
+                        help="The file to analyze (defaults to stdin if not provided)")
-    exit("File {} does not exist!".format(filename))
+    args = parser.parse_args()
-with open(filename, 'r') as fh:
-    data = json.loads(fh.read())
 
+    data = json.load(args.filename)
 
-# Go through data and disect every flow saved inside the dump
+    # Go through data and disect every flow saved inside the dump
-for export in sorted(data):
+    for key in sorted(data):
-    timestamp = datetime.fromtimestamp(float(export)).strftime("%Y-%m-%d %H:%M.%S")
+        timestamp = datetime.fromtimestamp(float(key)).strftime("%Y-%m-%d %H:%M.%S")
 
-    flows = data[export]
+        flows = data[key]
-    pending = None  # Two flows normally appear together for duplex connection
+        pending = None  # Two flows normally appear together for duplex connection
-    for flow in flows:
+        for flow in flows:
-        if not pending:
+            if not pending:
-            pending = flow
+                pending = flow
-        else:
+                continue
             con = Connection(pending, flow)
-            print("{timestamp}: {service:7} | {size:8} | {duration:9} | {src_host} ({src}) to"\
+            print("{timestamp}: {service:7} | {size:8} | {duration:9} | {src_host} ({src}) to {dest_host} ({dest})" \
-                " {dest_host} ({dest})".format(
+                .format(timestamp=timestamp, service=con.service.upper(), src_host=con.hostnames.src, src=con.src,
-                timestamp=timestamp, service=con.service.upper(),
+                        dest_host=con.hostnames.dest, dest=con.dest, size=con.human_size, duration=con.human_duration))
-                src_host=con.hostnames.src, src=con.src,
-                dest_host=con.hostnames.dest, dest=con.dest,
-                size=con.human_size, duration=con.human_duration))
             pending = None

main.py

@@ -1,140 +1,201 @@
 #!/usr/bin/env python3
 
 """
-Example collector script for NetFlow v9.
+Example collector script for NetFlow v1, v5, and v9.
 This file belongs to https://github.com/bitkeks/python-netflow-v9-softflowd.
 
 Copyright 2017-2019 Dominik Pataky <dev@bitkeks.eu>
 Licensed under MIT License. See LICENSE.
 """
 
-import logging
 import argparse
+from collections import namedtuple
+import queue
+import json
+import logging
 import sys
 import socketserver
+import threading
 import time
-import json
+
-import os.path
+from netflow import parse_packet, TemplateNotRecognized, UnknownNetFlowVersion
 
 
-logging.getLogger().setLevel(logging.INFO)
+logger = logging.getLogger(__name__)
-ch = logging.StreamHandler(sys.stdout)
-ch.setLevel(logging.DEBUG)
-formatter = logging.Formatter('%(message)s')
-ch.setFormatter(formatter)
-logging.getLogger().addHandler(ch)
 
-try:
+# Amount of time to wait before dropping an undecodable ExportPacket
-    from netflow.collector_v9 import ExportPacket, TemplateNotRecognized
+PACKET_TIMEOUT = 60 * 60
-except ImportError:
-    logging.warning("Netflow v9 not installed as package! Running from directory.")
-    from src.netflow.collector_v9 import ExportPacket, TemplateNotRecognized
 
-parser = argparse.ArgumentParser(description="A sample netflow collector.")
+# TODO: Add source IP
-parser.add_argument("--host", type=str, default="",
+RawPacket = namedtuple('RawPacket', ['ts', 'data'])
-                    help="collector listening address")
-parser.add_argument("--port", "-p", type=int, default=2055,
-                    help="collector listener port")
-parser.add_argument("--file", "-o", type=str, dest="output_file",
-                    default="{}.json".format(int(time.time())),
-                    help="collector export JSON file")
-parser.add_argument("--debug", "-D", action="store_true",
-                    help="Enable debug output")
 
 
-class SoftflowUDPHandler(socketserver.BaseRequestHandler):
+class QueuingRequestHandler(socketserver.BaseRequestHandler):
-    # We need to save the templates our NetFlow device
-    # send over time. Templates are not resended every
-    # time a flow is sent to the collector.
-    templates = {}
-    buffered = {}
-
-    @classmethod
-    def set_output_file(cls, path):
-        cls.output_file = path
-
     def handle(self):
-        if not os.path.exists(self.output_file):
-            with open(self.output_file, 'w') as fh:
-                json.dump({}, fh)
-
-        with open(self.output_file, 'r') as fh:
-            try:
-                existing_data = json.load(fh)
-            except json.decoder.JSONDecodeError as ex:
-                logging.error("Malformed JSON output file. Cannot read existing data, aborting.")
-                return
-
         data = self.request[0]
-        host = self.client_address[0]
+        self.server.queue.put(RawPacket(time.time(), data))
-        logging.debug("Received data from {}, length {}".format(host, len(data)))
+        logger.debug(
+            "Received %d bytes of data from %s", len(data), self.client_address[0]
+        )
 
-        export = None
+
+class QueuingUDPListener(socketserver.ThreadingUDPServer):
+    """A threaded UDP server that adds a (time, data) tuple to a queue for
+    every request it sees
+    """
+    def __init__(self, interface, queue):
+        self.queue = queue
+        super().__init__(interface, QueuingRequestHandler)
+
+
+class NetFlowListener(threading.Thread):
+    """A thread that listens for incoming NetFlow packets, processes them, and
+    makes them available to consumers.
+
+    - When initialized, will start listening for NetFlow packets on the provided
+      host and port and queuing them for processing.
+    - When started, will start processing and parsing queued packets.
+    - When stopped, will shut down the listener and stop processing.
+    - When joined, will wait for the listener to exit
+
+    For example, a simple script that outputs data until killed with CTRL+C:
+    >>> listener = NetFlowListener('0.0.0.0', 2055)
+    >>> print("Listening for NetFlow packets")
+    >>> listener.start() # start processing packets
+    >>> try:
+    ...     while True:
+    ...         ts, export = listener.get()
+    ...         print("Time: {}".format(ts))
+    ...         for f in export.flows:
+    ...             print(" - {IPV4_SRC_ADDR} sent data to {IPV4_DST_ADDR}"
+    ...                   "".format(**f))
+    ... finally:
+    ...     print("Stopping...")
+    ...     listener.stop()
+    ...     listener.join()
+    ...     print("Stopped!")
+    """
+
+    def __init__(self, host, port):
+        logger.info("Starting the NetFlow listener on {}:{}".format(host, port))
+        self.output = queue.Queue()
+        self.input = queue.Queue()
+        self.server = QueuingUDPListener((host, port), self.input)
+        self.thread = threading.Thread(target=self.server.serve_forever)
+        self.thread.start()
+        self._shutdown = threading.Event()
+        super().__init__()
+
+    def get(self, block=True, timeout=None):
+        """Get a processed flow.
+
+        If optional args 'block' is true and 'timeout' is None (the default),
+        block if necessary until a flow is available. If 'timeout' is
+        a non-negative number, it blocks at most 'timeout' seconds and raises
+        the queue.Empty exception if no flow was available within that time.
+        Otherwise ('block' is false), return a flow if one is immediately
+        available, else raise the queue.Empty exception ('timeout' is ignored
+        in that case).
+        """
+        return self.output.get(block, timeout)
+
+    def run(self):
+        # Process packets from the queue
         try:
-            export = ExportPacket(data, self.templates)
+            templates = {}
-        except TemplateNotRecognized:
+            to_retry = []
-            self.buffered[time.time()] = data
+            while not self._shutdown.is_set():
-            logging.warning("Received data with unknown template, data stored in buffer!")
+                try:
-            return
+                    # 0.5s delay to limit CPU usage while waiting for new packets
+                    pkt = self.input.get(block=True, timeout=0.5)
+                except queue.Empty:
+                    continue
 
-        if not export:
+                try:
-            logging.error("Error with exception handling while disecting export, export is None")
+                    export = parse_packet(pkt.data, templates)
-            return
+                except UnknownNetFlowVersion as e:
+                    logger.error("%s, ignoring the packet", e)
+                    continue
+                except TemplateNotRecognized:
+                    if time.time() - pkt.ts > PACKET_TIMEOUT:
+                        logger.warning("Dropping an old and undecodable v9 ExportPacket")
+                    else:
+                        to_retry.append(pkt)
+                        logger.debug("Failed to decode a v9 ExportPacket - will "
+                                      "re-attempt when a new template is discovered")
+                    continue
 
-        logging.debug("Processed ExportPacket with {} flows.".format(export.header.count))
+                logger.debug("Processed a v%d ExportPacket with %d flows.",
-        logging.debug("Size of buffer: {}".format(len(self.buffered)))
+                             export.header.version, export.header.count)
 
-        # In case the export held some new templates
+                # If any new templates were discovered, dump the unprocessable
-        self.templates.update(export.templates)
+                # data back into the queue and try to decode them again
+                if (export.header.version == 9 and export.contains_new_templates and to_retry):
+                    logger.debug("Received new template(s)")
+                    logger.debug("Will re-attempt to decode %d old v9 ExportPackets",
+                                 len(to_retry))
+                    for p in to_retry:
+                        self.input.put(p)
+                    to_retry.clear()
 
-        remain_buffered = {}
+                self.output.put((pkt.ts, export))
-        processed = []
+        finally:
-        for timestamp, data in self.buffered.items():
+            self.server.shutdown()
-            try:
+            self.server.server_close()
-                buffered_export = ExportPacket(data, self.templates)
-                processed.append(timestamp)
-            except TemplateNotRecognized:
-                remain_buffered[timestamp] = data
-                logging.debug("Template of buffered ExportPacket still not recognized")
-                continue
-            logging.debug("Processed buffered ExportPacket with {} flows.".format(buffered_export.header.count))
-            existing_data[timestamp] = [flow.data for flow in buffered_export.flows]
 
-        # Delete processed items from the buffer
+    def stop(self):
-        for pro in processed:
+        logger.info("Shutting down the NetFlow listener")
-            del self.buffered[pro]
+        self._shutdown.set()
 
-        # Update the buffer
+    def join(self, timeout=None):
-        self.buffered.update(remain_buffered)
+        self.thread.join(timeout=timeout)
+        super().join(timeout=timeout)
 
-        # Append new flows
-        existing_data[time.time()] = [flow.data for flow in export.flows]
 
-        with open(self.output_file, 'w') as fh:
+def get_export_packets(host, port):
-            json.dump(existing_data, fh)
+    """A generator that will yield ExportPacket objects until it is killed"""
 
+    listener = NetFlowListener(host, port)
+    listener.start()
+    try:
+        while True:
+            yield listener.get()
+    finally:
+        listener.stop()
+        listener.join()
 
 
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="A sample netflow collector.")
+    parser.add_argument("--host", type=str, default="0.0.0.0",
+                        help="collector listening address")
+    parser.add_argument("--port", "-p", type=int, default=2055,
+                        help="collector listener port")
+    parser.add_argument("--file", "-o", type=str, dest="output_file",
+                        default="{}.json".format(int(time.time())),
+                        help="collector export JSON file")
+    parser.add_argument("--debug", "-D", action="store_true",
+                        help="Enable debug output")
     args = parser.parse_args()
 
+    logging.basicConfig(level=logging.INFO, stream=sys.stdout, format="%(message)s")
+
     if args.debug:
-        logging.getLogger().setLevel(logging.DEBUG)
+        logger.setLevel(logging.DEBUG)
-
-    output_file = args.output_file
-    SoftflowUDPHandler.set_output_file(output_file)
-
-    host = args.host
-    port = args.port
-    logging.info("Listening on interface {}:{}".format(host, port))
-    server = socketserver.UDPServer((host, port), SoftflowUDPHandler)
 
+    data = {}
     try:
-        logging.debug("Starting the NetFlow listener")
+        # TODO: For a long-running processes, this will consume loads of memory
-        server.serve_forever(poll_interval=0.5)
+        for ts, export in get_export_packets(args.host, args.port):
-    except (IOError, SystemExit):
+            data[ts] = [flow.data for flow in export.flows]
-        raise
     except KeyboardInterrupt:
-        raise
+        logger.info("Received KeyboardInterrupt, passing through")
+        pass
 
-    server.server_close()
+    if data:
+        # TODO: this should be done periodically to not lose any data (only saved in memory)
+        logger.info("Outputting collected data to '%s'", args.output_file)
+        with open(args.output_file, 'w') as f:
+            json.dump(data, f)
+    else:
+        logger.info("No data collected")

netflow/__init__.py

@@ -0,0 +1,35 @@
+#!/usr/bin/env python
+
+import struct
+
+from netflow.v1 import V1ExportPacket
+from netflow.v5 import V5ExportPacket
+from netflow.v9 import V9ExportPacket, TemplateNotRecognized
+
+__all__ = ["TemplateNotRecognized", "UnknownNetFlowVersion", "parse_packet"]
+
+
+class UnknownNetFlowVersion(Exception):
+    def __init__(self, data, version):
+        self.data = data
+        self.version = version
+        r = repr(data)
+        data_str = ("{:.25}..." if len(r) >= 28 else "{}").format(r)
+        super().__init__(
+            "Unknown NetFlow version {} for data {}".format(version, data_str)
+        )
+
+
+def get_netflow_version(data):
+    return struct.unpack('!H', data[:2])[0]
+
+
+def parse_packet(data, templates):
+    version = get_netflow_version(data)
+    if version == 1:
+        return V1ExportPacket(data)
+    elif version == 5:
+        return V5ExportPacket(data)
+    elif version == 9:
+        return V9ExportPacket(data, templates)
+    raise UnknownNetFlowVersion(data, version)

netflow/v1.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python3
+
+"""
+Netflow V1 collector and parser implementation in Python 3.
+Created purely for fun. Not battled tested nor will it be.
+
+Reference https://www.cisco.com/c/en/us/td/docs/net_mgmt/netflow_collection_engine/3-6/user/guide/format.html
+
+This script is specifically implemented in combination with softflowd.
+See https://github.com/djmdjm/softflowd
+
+"""
+
+import struct
+
+
+class DataFlow:
+    """Holds one v1 DataRecord"""
+
+    length = 48
+
+    def __init__(self, data):
+        self.data = {}
+        self.data['IPV4_SRC_ADDR'] = struct.unpack('!I', data[:4])[0]
+        self.data['IPV4_DST_ADDR'] = struct.unpack('!I', data[4:8])[0]
+        self.data['NEXT_HOP'] = struct.unpack('!I', data[8:12])[0]
+        self.data['INPUT'] = struct.unpack('!H', data[12:14])[0]
+        self.data['OUTPUT'] = struct.unpack('!H', data[14:16])[0]
+        self.data['IN_PACKETS'] = struct.unpack('!I', data[16:20])[0]
+        self.data['IN_OCTETS'] = struct.unpack('!I', data[20:24])[0]
+        self.data['FIRST_SWITCHED'] = struct.unpack('!I', data[24:28])[0]
+        self.data['LAST_SWITCHED'] = struct.unpack('!I', data[28:32])[0]
+        self.data['SRC_PORT'] = struct.unpack('!H', data[32:34])[0]
+        self.data['DST_PORT'] = struct.unpack('!H', data[34:36])[0]
+        # Word at 36 is used for padding
+        self.data['PROTO'] = struct.unpack('!B', data[38:39])[0]
+        self.data['TOS'] = struct.unpack('!B', data[39:40])[0]
+        self.data['TCP_FLAGS'] = struct.unpack('!B', data[40:41])[0]
+        # Data at 41-48 is padding
+
+    def __repr__(self):
+        return "<DataRecord with data {}>".format(self.data)
+
+
+class Header:
+    """The header of the V1ExportPacket"""
+
+    length = 16
+
+    def __init__(self, data):
+        header = struct.unpack('!HHIII', data[:self.length])
+        self.version = header[0]
+        self.count = header[1]
+        self.uptime = header[2]
+        self.timestamp = header[3]
+        self.timestamp_nano = header[4]
+
+
+class V1ExportPacket:
+    """The flow record holds the header and data flowsets."""
+
+    def __init__(self, data):
+        self.flows = []
+        self.header = Header(data)
+
+        offset = self.header.length
+        for flow_count in range(0, self.header.count):
+            flow = V1DataFlow(data[offset:])
+            self.flows.append(flow)
+            offset += flow.length
+
+    def __repr__(self):
+        return "<ExportPacket v{} with {} records>".format(
+                self.header.version, self.header.count)

netflow/v5.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python3
+
+"""
+Netflow V5 collector and parser implementation in Python 3.
+Created purely for fun. Not battled tested nor will it be.
+
+Reference: https://www.cisco.com/c/en/us/td/docs/net_mgmt/netflow_collection_engine/3-6/user/guide/format.html
+
+This script is specifically implemented in combination with softflowd.
+See https://github.com/djmdjm/softflowd
+
+"""
+
+import struct
+
+
+class DataFlow:
+    """Holds one v5 DataRecord"""
+
+    length = 48
+
+    def __init__(self, data):
+        self.data = {}
+        self.data['IPV4_SRC_ADDR'] = struct.unpack('!I', data[:4])[0]
+        self.data['IPV4_DST_ADDR'] = struct.unpack('!I', data[4:8])[0]
+        self.data['NEXT_HOP'] = struct.unpack('!I', data[8:12])[0]
+        self.data['INPUT'] = struct.unpack('!H', data[12:14])[0]
+        self.data['OUTPUT'] = struct.unpack('!H', data[14:16])[0]
+        self.data['IN_PACKETS'] = struct.unpack('!I', data[16:20])[0]
+        self.data['IN_OCTETS'] = struct.unpack('!I', data[20:24])[0]
+        self.data['FIRST_SWITCHED'] = struct.unpack('!I', data[24:28])[0]
+        self.data['LAST_SWITCHED'] = struct.unpack('!I', data[28:32])[0]
+        self.data['SRC_PORT'] = struct.unpack('!H', data[32:34])[0]
+        self.data['DST_PORT'] = struct.unpack('!H', data[34:36])[0]
+        # Byte 36 is used for padding
+        self.data['TCP_FLAGS'] = struct.unpack('!B', data[37:38])[0]
+        self.data['PROTO'] = struct.unpack('!B', data[38:39])[0]
+        self.data['TOS'] = struct.unpack('!B', data[39:40])[0]
+        self.data['SRC_AS'] = struct.unpack('!H', data[40:42])[0]
+        self.data['DST_AS'] = struct.unpack('!H', data[42:44])[0]
+        self.data['SRC_MASK'] = struct.unpack('!B', data[44:45])[0]
+        self.data['DST_MASK'] = struct.unpack('!B', data[45:46])[0]
+        # Word 46 is used for padding
+
+    def __repr__(self):
+        return "<DataRecord with data {}>".format(self.data)
+
+
+class Header:
+    """The header of the V5ExportPacket"""
+
+    length = 24
+
+    def __init__(self, data):
+        header = struct.unpack('!HHIIIIBBH', data[:self.length])
+        self.version = header[0]
+        self.count = header[1]
+        self.uptime = header[2]
+        self.timestamp = header[3]
+        self.timestamp_nano = header[4]
+        self.sequence = header[5]
+        self.engine_type = header[6]
+        self.engine_id = header[7]
+        self.sampling_interval = header[8]
+
+
+class V5ExportPacket:
+    """The flow record holds the header and data flowsets."""
+
+    def __init__(self, data):
+        self.flows = []
+        self.header = Header(data)
+
+        offset = self.header.length
+        for flow_count in range(0, self.header.count):
+            flow = DataFlow(data[offset:])
+            self.flows.append(flow)
+            offset += flow.length
+
+    def __repr__(self):
+        return "<ExportPacket v{} with {} records>".format(
+                self.header.version, self.header.count)

netflow/v9.py

@@ -4,6 +4,8 @@
 Netflow V9 collector and parser implementation in Python 3.
 Created for learning purposes and unsatisfying alternatives.
 
+Reference: https://www.cisco.com/en/US/technologies/tk648/tk362/technologies_white_paper09186a00800a3db9.html
+
 This script is specifically implemented in combination with softflowd.
 See https://github.com/djmdjm/softflowd
 
@@ -11,12 +13,10 @@ Copyright 2017, 2018 Dominik Pataky <dev@bitkeks.eu>
 Licensed under MIT License. See LICENSE.
 """
 
-import socket
 import struct
-import sys
 
 
-field_types = {
+FIELD_TYPES = {
     0: 'UNKNOWN_FIELD_TYPE',  # fallback for unknown field types
 
     # Cisco specs for NetFlow v9
@@ -153,10 +153,14 @@ field_types = {
 }
 
 
+class TemplateNotRecognized(KeyError):
+    pass
+
+
 class DataRecord:
     """This is a 'flow' as we want it from our source. What it contains is
     variable in NetFlow V9, so to work with the data you have to analyze the
-    data dict keys (which are integers and can be mapped with the field_types
+    data dict keys (which are integers and can be mapped with the FIELD_TYPES
     dict).
 
     Should hold a 'data' dict with keys=field_type (integer) and value (in bytes).
@@ -195,7 +199,7 @@ class DataFlowSet:
 
             for field in template.fields:
                 flen = field.field_length
-                fkey = field_types[field.field_type]
+                fkey = FIELD_TYPES[field.field_type]
                 fdata = None
 
                 # The length of the value byte slice is defined in the template
@@ -218,20 +222,18 @@ class DataFlowSet:
 
 
 class TemplateField:
-    """A field with type identifier and length.
+    """A field with type identifier and length."""
-    """
     def __init__(self, field_type, field_length):
         self.field_type = field_type  # integer
         self.field_length = field_length  # bytes
 
     def __repr__(self):
         return "<TemplateField type {}:{}, length {}>".format(
-            self.field_type, field_types[self.field_type], self.field_length)
+            self.field_type, FIELD_TYPES[self.field_type], self.field_length)
 
 
 class TemplateRecord:
-    """A template record contained in a TemplateFlowSet.
+    """A template record contained in a TemplateFlowSet."""
-    """
     def __init__(self, template_id, field_count, fields):
         self.template_id = template_id
         self.field_count = field_count
@@ -240,7 +242,7 @@ class TemplateRecord:
     def __repr__(self):
         return "<TemplateRecord {} with {} fields: {}>".format(
             self.template_id, self.field_count,
-            ' '.join([field_types[field.field_type] for field in self.fields]))
+            ' '.join([FIELD_TYPES[field.field_type] for field in self.fields]))
 
 
 class TemplateFlowSet:
@@ -258,7 +260,7 @@ class TemplateFlowSet:
         offset = 4  # Skip header
 
         # Iterate through all template records in this template flowset
-        while offset != self.length:
+        while offset < self.length:
             pack = struct.unpack('!HH', data[offset:offset+4])
             template_id = pack[0]
             field_count = pack[1]
@@ -268,7 +270,7 @@ class TemplateFlowSet:
                 # Get all fields of this template
                 offset += 4
                 field_type, field_length = struct.unpack('!HH', data[offset:offset+4])
-                if field_type not in field_types:
+                if field_type not in FIELD_TYPES:
                     field_type = 0  # Set field_type to UNKNOWN_FIELD_TYPE as fallback
                 field = TemplateField(field_type, field_length)
                 fields.append(field)
@@ -288,10 +290,12 @@ class TemplateFlowSet:
 
 
 class Header:
-    """The header of the ExportPacket.
+    """The header of the V9ExportPacket"""
-    """
+
+    length = 20
+
     def __init__(self, data):
-        pack = struct.unpack('!HHIIII', data[:20])
+        pack = struct.unpack('!HHIIII', data[:self.length])
 
         self.version = pack[0]
         self.count = pack[1]  # not sure if correct. softflowd: no of flows
@@ -301,19 +305,26 @@ class Header:
         self.source_id = pack[5]
 
 
-class ExportPacket:
+class V9ExportPacket:
-    """The flow record holds the header and all template and data flowsets.
+    """The flow record holds the header and all template and data flowsets."""
-    """
+
     def __init__(self, data, templates):
         self.header = Header(data)
         self.templates = templates
+        self._new_templates = False
         self.flows = []
 
-        offset = 20
+        offset = self.header.length
         while offset != len(data):
             flowset_id = struct.unpack('!H', data[offset:offset+2])[0]
             if flowset_id == 0:  # TemplateFlowSet always have id 0
                 tfs = TemplateFlowSet(data[offset:])
+                # Check for any new/changed templates
+                if not self._new_templates:
+                    for id_, template in tfs.templates.items():
+                        if id_ not in self.templates or self.templates[id_] != template:
+                            self._new_templates = True
+                            break
                 self.templates.update(tfs.templates)
                 offset += tfs.length
             else:
@@ -321,10 +332,11 @@ class ExportPacket:
                 self.flows += dfs.flows
                 offset += dfs.length
 
+    @property
+    def contains_new_templates(self):
+        return self._new_templates
+
     def __repr__(self):
-        return "<ExportPacket version {} counting {} records>".format(
+        s = " and new template(s)" if self.contains_new_templates else ""
-            self.header.version, self.header.count)
+        return "<ExportPacket v{} with {} records{}>".format(
-
+            self.header.version, self.header.count, s)
-
-class TemplateNotRecognized(KeyError):
-    pass

setup.py

@@ -1,17 +1,12 @@
 #!/usr/bin/env python3
 
-from setuptools import setup, find_packages
+from setuptools import setup
-import os
 
-data_files = [(d, [os.path.join(d, f) for f in files])
+setup(name='netflow',
-              for d, folders, files in os.walk(os.path.join('src', 'config'))]
+      version='0.8.0',
-
+      description='NetFlow v1, v5, and v9 collector, parser and analyzer implemented in Python 3.',
-setup(name='netflow-v9',
-      version='0.7.0',
-      description='NetFlow v9 parser and collector implemented in Python 3. Developed to be used with softflowd v0.9.9',
       author='Dominik Pataky',
       author_email='dev@bitkeks.eu',
-      packages=find_packages('src'),
+      packages=["netflow"],
-      package_dir={'': 'src'},
       license='MIT'
 )

tests.py

@@ -1,144 +1,150 @@
 #!/usr/bin/env python3
 
 """
-This file contains tests for the softflowd UDP collector saved in main.py
+This file contains tests for the softflowd UDP collector saved in main.py The
-The test packets (defined below as hex streams) were extracted from a "real" softflowd export
+test packets (defined below as hex streams) were extracted from a "real"
-based on a sample PCAP capture file. They consist of one export with the templates and three without.
+softflowd export based on a sample PCAP capture file. They consist of one
-
+export with the templates and three without.
-Two tests are defined, one slow, one fast. During some runs exceptions occured which might hint
-to race conditions during reading and writing to the JSON output file.
-For now, both tests run successfully.
 
 Copyright 2017-2019 Dominik Pataky <dev@bitkeks.eu>
 Licensed under MIT License. See LICENSE.
 """
 
-import ipaddress
 import json
 import logging
-from pprint import pprint
+import queue
 import random
 import socket
-import socketserver
 import subprocess
-import tempfile
+import sys
-from time import sleep
+import time
-import threading
 import unittest
 
-from main import SoftflowUDPHandler
+from main import NetFlowListener
 
-
+# TODO: add tests for v1 and v5
-logging.getLogger().setLevel(logging.DEBUG)
+# TODO: tests with 500 packets fail?
 
 # The flowset with 2 templates and 8 flows
-template_packet = '0009000a000000035c9f55980000000100000000000000400400000e00080004000c000400150004001600040001000400020004000a0004000e000400070002000b00020004000100060001003c000100050001000000400800000e001b0010001c001000150004001600040001000400020004000a0004000e000400070002000b00020004000100060001003c000100050001040001447f0000017f000001fb3c1aaafb3c18fd000190100000004b00000000000000000050942c061b04007f0000017f000001fb3c1aaafb3c18fd00000f94000000360000000000000000942c0050061f04007f0000017f000001fb3c1cfcfb3c1a9b0000d3fc0000002a000000000000000000509434061b04007f0000017f000001fb3c1cfcfb3c1a9b00000a490000001e000000000000000094340050061f04007f0000017f000001fb3bb82cfb3ba48b000002960000000300000000000000000050942a061904007f0000017f000001fb3bb82cfb3ba48b00000068000000020000000000000000942a0050061104007f0000017f000001fb3c1900fb3c18fe0000004c0000000100000000000000000035b3c9110004007f0000017f000001fb3c1900fb3c18fe0000003c000000010000000000000000b3c9003511000400'
+TEMPLATE_PACKET = '0009000a000000035c9f55980000000100000000000000400400000e00080004000c000400150004001600040001000400020004000a0004000e000400070002000b00020004000100060001003c000100050001000000400800000e001b0010001c001000150004001600040001000400020004000a0004000e000400070002000b00020004000100060001003c000100050001040001447f0000017f000001fb3c1aaafb3c18fd000190100000004b00000000000000000050942c061b04007f0000017f000001fb3c1aaafb3c18fd00000f94000000360000000000000000942c0050061f04007f0000017f000001fb3c1cfcfb3c1a9b0000d3fc0000002a000000000000000000509434061b04007f0000017f000001fb3c1cfcfb3c1a9b00000a490000001e000000000000000094340050061f04007f0000017f000001fb3bb82cfb3ba48b000002960000000300000000000000000050942a061904007f0000017f000001fb3bb82cfb3ba48b00000068000000020000000000000000942a0050061104007f0000017f000001fb3c1900fb3c18fe0000004c0000000100000000000000000035b3c9110004007f0000017f000001fb3c1900fb3c18fe0000003c000000010000000000000000b3c9003511000400'
 
 # Three packets without templates, each with 12 flows, anonymized
-packets = [
+PACKETS = [
     '0009000c000000035c9f55980000000200000000040001e47f0000017f000001fb3c1a17fb3c19fd000001480000000200000000000000000035ea82110004007f0000017f000001fb3c1a17fb3c19fd0000007a000000020000000000000000ea820035110004007f0000017f000001fb3c1a17fb3c19fd000000f80000000200000000000000000035c6e2110004007f0000017f000001fb3c1a17fb3c19fd0000007a000000020000000000000000c6e20035110004007f0000017f000001fb3c1a9efb3c1a9c0000004c0000000100000000000000000035adc1110004007f0000017f000001fb3c1a9efb3c1a9c0000003c000000010000000000000000adc10035110004007f0000017f000001fb3c1b74fb3c1b720000004c0000000100000000000000000035d0b3110004007f0000017f000001fb3c1b74fb3c1b720000003c000000010000000000000000d0b30035110004007f0000017f000001fb3c2f59fb3c1b7100001a350000000a000000000000000000509436061b04007f0000017f000001fb3c2f59fb3c1b710000038a0000000a000000000000000094360050061b04007f0000017f000001fb3c913bfb3c91380000004c0000000100000000000000000035e262110004007f0000017f000001fb3c913bfb3c91380000003c000000010000000000000000e262003511000400',
     '0009000c000000035c9f55980000000300000000040001e47f0000017f000001fb3ca523fb3c913b0000030700000005000000000000000000509438061b04007f0000017f000001fb3ca523fb3c913b000002a200000005000000000000000094380050061b04007f0000017f000001fb3f7fe1fb3dbc970002d52800000097000000000000000001bb8730061b04007f0000017f000001fb3f7fe1fb3dbc970000146c000000520000000000000000873001bb061f04007f0000017f000001fb3d066ffb3d066c0000004c0000000100000000000000000035e5bd110004007f0000017f000001fb3d066ffb3d066c0000003c000000010000000000000000e5bd0035110004007f0000017f000001fb3d1a61fb3d066b000003060000000500000000000000000050943a061b04007f0000017f000001fb3d1a61fb3d066b000002a2000000050000000000000000943a0050061b04007f0000017f000001fb3fed00fb3f002c0000344000000016000000000000000001bbae50061f04007f0000017f000001fb3fed00fb3f002c00000a47000000120000000000000000ae5001bb061b04007f0000017f000001fb402f17fb402a750003524c000000a5000000000000000001bbc48c061b04007f0000017f000001fb402f17fb402a75000020a60000007e0000000000000000c48c01bb061f0400',
     '0009000c000000035c9f55980000000400000000040001e47f0000017f000001fb3d7ba2fb3d7ba00000004c0000000100000000000000000035a399110004007f0000017f000001fb3d7ba2fb3d7ba00000003c000000010000000000000000a3990035110004007f0000017f000001fb3d8f85fb3d7b9f000003070000000500000000000000000050943c061b04007f0000017f000001fb3d8f85fb3d7b9f000002a2000000050000000000000000943c0050061b04007f0000017f000001fb3d9165fb3d7f6d0000c97b0000002a000000000000000001bbae48061b04007f0000017f000001fb3d9165fb3d7f6d000007f40000001a0000000000000000ae4801bb061b04007f0000017f000001fb3dbc96fb3dbc7e0000011e0000000200000000000000000035bd4f110004007f0000017f000001fb3dbc96fb3dbc7e0000008e000000020000000000000000bd4f0035110004007f0000017f000001fb3ddbb3fb3c1a180000bfee0000002f00000000000000000050ae56061b04007f0000017f000001fb3ddbb3fb3c1a1800000982000000270000000000000000ae560050061b04007f0000017f000001fb3ddbb3fb3c1a180000130e0000001200000000000000000050e820061b04007f0000017f000001fb3ddbb3fb3c1a180000059c000000140000000000000000e8200050061b0400'
 ]
 
+INVALID_PACKET = "FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF"
 
-class ThreadedUDPServer(socketserver.ThreadingMixIn, socketserver.UDPServer):
+CONNECTION = ('127.0.0.1', 1337)
-    pass
+NUM_PACKETS = 50
 
-class TestSoftflowExport(unittest.TestCase):
-    CONNECTION = ('127.0.0.1', 1337)
-    COUNT_PACKETS_TO_TEST = 5
-    SLEEP_TIME = 0.3
-    RUN_ANALYZER = False
 
-    def setUp(self):
+def emit_packets(packets, delay=0):
-        logging.debug("Creating temporary JSON output file.")
+    """Send the provided packets to the listener"""
-        self.temp_output_file = tempfile.NamedTemporaryFile(prefix="softflowd_")
+    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    for p in packets:
+        sock.sendto(bytes.fromhex(p), CONNECTION)
+        time.sleep(delay)
+    sock.close()
 
-        # FIXME: templates are saved between test runs, because they are stored with the class
-        # Maybe the templates should be stored with an instance?
-        logging.debug("Resetting SoftflowUDPHandler templates.")
-        SoftflowUDPHandler.templates = {}
 
-        logging.debug("Setting temporary file {} as output for SoftflowUDPHandler".format(self.temp_output_file.name))
+def send_recv_packets(packets, delay=0):
-        SoftflowUDPHandler.set_output_file(self.temp_output_file.name)
+    """Starts a listener, send packets, receives packets
 
-        logging.debug("Writing empty dict to output file.")
+    returns a tuple: ([(ts, export), ...], time_started_sending, time_stopped_sending)
-        with open(self.temp_output_file.name, "w") as fh:
+    """
-            json.dump({}, fh)
+    l = NetFlowListener(*CONNECTION)
+    tstart = time.time()
+    emit_packets(packets, delay=delay)
+    time.sleep(0.5) # Allow packets to be sent and recieved
+    tend = time.time()
+    l.start()
 
-        logging.debug("Creating and running the Softflow collector in another thread.")
+    pkts = []
-        self.server = ThreadedUDPServer(self.CONNECTION, SoftflowUDPHandler)
+    while True:
-        self.server_thread = threading.Thread(target=self.server.serve_forever)
+        try:
-        self.server_thread.daemon = True
+            pkts.append(l.get(timeout=0.5))
-        self.server_thread.start()
+        except queue.Empty:
+            break
+    l.stop()
+    l.join()
+    return pkts, tstart, tend
 
-        logging.debug("Creating UDP socket for client packets.")
-        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
 
-    def tearDown(self):
+class TestSoftFlowExport(unittest.TestCase):
-        logging.debug("Running tear down procedure.")
-        self.server.shutdown()
-        self.server.server_close()
-        self.server_thread.join()
-        self.sock.close()
-        self.temp_output_file.close()
 
-    def _test_export(self):
+    def _test_recv_all_packets(self, num, template_idx, delay=0):
-        logging.info("Running UDP client sending raw hex packets with flows.")
+        """Fling packets at the server and test that it receives them all"""
+        def gen_pkts(n, idx):
+            for x in range(n):
+                if x == idx:
+                    yield TEMPLATE_PACKET
+                else:
+                    yield random.choice(PACKETS)
 
-        # Get a random index on which the template is sent
+        pkts, tstart, tend = send_recv_packets(gen_pkts(num, template_idx), delay=delay)
-        template_idx = random.randint(1, self.COUNT_PACKETS_TO_TEST - 1)  # 1 for enhanced testing, -1 because randint
 
-        # Save the order of lengths for later check
+        # check number of packets
-        lens = []
+        self.assertEqual(len(pkts), num)
 
-        for idx in range(self.COUNT_PACKETS_TO_TEST):
+        # check timestamps are when packets were sent, not processed
-            # Choose a random packet payload
+        self.assertTrue(all(tstart < p[0] < tend for p in pkts))
-            p = random.choice(packets)
 
-            logging.info("Sending packet {}.".format(idx))
+        # check number of "things" in the packets (flows + templates)
-            self.sock.sendto(bytes.fromhex(p), self.CONNECTION)
+        # template packet = 10 things
-            lens.append(12)
+        # other packets = 12 things
-            sleep(self.SLEEP_TIME)
+        self.assertEqual(sum(p[1].header.count for p in pkts), (num-1)*12 + 10)
 
-            # Randomly inject the template packet
+        # check number of flows in the packets
-            if idx == template_idx:
+        # template packet = 8 flows (2 templates)
-                logging.info("Sending template packet.")
+        # other packets = 12 flows
-                self.sock.sendto(bytes.fromhex(template_packet), self.CONNECTION)
+        self.assertEqual(sum(len(p[1].flows) for p in pkts), (num-1)*12 + 8)
-                lens.append(8)
-                sleep(self.SLEEP_TIME)
 
-        with open(self.temp_output_file.name, "r") as fh:
+    def test_recv_all_packets_template_first(self):
-            exported = json.load(fh)
+        """Test all packets are received when the template is sent first"""
+        self._test_recv_all_packets(NUM_PACKETS, 0)
 
-        # We got four exports
+    def test_recv_all_packets_template_middle(self):
-        logging.info("Testing the existence of all exports, including the ones with formerly unknown templates: {} of {}".format(
+        """Test all packets are received when the template is sent in the middle"""
-            len(exported.keys()), self.COUNT_PACKETS_TO_TEST + 1))
+        self._test_recv_all_packets(NUM_PACKETS, NUM_PACKETS//2)
-        self.assertEqual(len(exported.keys()), self.COUNT_PACKETS_TO_TEST + 1)  # +1 including the template packet
 
-        # Test lengths of exports
+    def test_recv_all_packets_template_last(self):
-        logging.info("Testing the correct lengths of all exports.")
+        """Test all packets are received when the template is sent last"""
-        for idx, val in enumerate(exported.values()):
+        self._test_recv_all_packets(NUM_PACKETS, NUM_PACKETS-1)
-            self.assertEqual(len(val), lens[idx])
 
-        if self.RUN_ANALYZER:
+    def test_recv_all_packets_slowly(self):
-            logging.info("Running analyze_json.py")
+        """Test all packets are received when things are sent slooooowwwwwwwwlllllllyyyyyy"""
-            analyzer = subprocess.run(['python3', 'analyze_json.py', self.temp_output_file.name], stdout=subprocess.PIPE)
+        self._test_recv_all_packets(3, 0, delay=1)
-            for line in analyzer.stdout.split(b"\n"):
-                print(line.decode())
 
-    def test_slow(self):
+    def test_ignore_invalid_packets(self):
-        logging.info("Running slow test")
+        """Test that invlalid packets log a warning but are otherwise ignored"""
-        self.SLEEP_TIME = 0.5
+        with self.assertLogs(level='WARNING'):
-        self.COUNT_PACKETS_TO_TEST = 3
+            pkts, _, _ = send_recv_packets([
-        self._test_export()
+                INVALID_PACKET, TEMPLATE_PACKET, random.choice(PACKETS), INVALID_PACKET,
+                random.choice(PACKETS), INVALID_PACKET
+            ])
+        self.assertEqual(len(pkts), 3)
+
+    def test_analyzer(self):
+        """Test thar the analyzer doesn't break and outputs the correct number of lines"""
+        pkts, _, _ = send_recv_packets([TEMPLATE_PACKET, *PACKETS])
+        data = {p[0]: [f.data for f in p[1].flows] for p in pkts}
+        analyzer = subprocess.run(
+            [sys.executable, 'analyze_json.py'],
+            input=json.dumps(data),
+            encoding='utf-8',
+            capture_output=True
+        )
+
+        # every 2 flows are written as a single line (any extras are dropped)
+        num_flows = sum(len(f) for f in data.values())
+        self.assertEqual(len(analyzer.stdout.splitlines()), num_flows//2)
+
+        # make sure there are no errors
+        self.assertEqual(analyzer.stderr, "")
 
-    def test_fast(self):
-        logging.info("Running fast test")
-        self.SLEEP_TIME = 0.1
-        self.COUNT_PACKETS_TO_TEST = 30
-        self._test_export()
 
 if __name__ == '__main__':
     unittest.main()