Implement NetFlow bot

Author: grafolean

lookup.py

@@ -1,3 +1,5 @@
+REDIS_HASH_TRAFFIC_PER_PROTOCOL = "tpp"
+
 PROTOCOLS = {
     0: "HOPOPT",
     1: "ICMP",

netflowbot.py

@@ -10,9 +10,11 @@
 from colors import color
 import dotenv
 import redis
+import requests
 
-from grafoleancollector import Collector
+from grafoleancollector import Collector, send_results_to_grafolean
 
+from lookup import PROTOCOLS, REDIS_HASH_TRAFFIC_PER_PROTOCOL
 
 logging.basicConfig(format='%(asctime)s | %(levelname)s | %(message)s',
                     datefmt='%Y-%m-%d %H:%M:%S', level=logging.DEBUG)
@@ -27,8 +29,95 @@
 r = redis.Redis(host=REDIS_HOST)
 
 
+REDIS_PREFIX = 'netflow'
+
+
 class NetFlowBot(Collector):
-    pass
+
+    def jobs(self):
+        # for entity_info in self.fetch_job_configs('netflow'):
+        #     for sensor_info in entity_info["sensors"]:
+        #         # The job could be triggered at different intervals - it is triggered when at least one of the specified intervals matches.
+        #         intervals = [sensor_info["interval"]]
+        #         # `job_id` must be a unique, permanent identifier of a job. When the job_id changes, the job will be rescheduled - so make sure it is something that
+        #         # identifies this particular job.
+        #         job_id = str(sensor_info["sensor_id"])
+        #         # Prepare parameters that will be passed to `perform_job()` whenever the job is being run:
+        #         # (don't forget to pass backend_url and bot_token!)
+        #         job_params = { **sensor_info, "entity_info": entity_info, "backend_url": self.backend_url, "bot_token": self.bot_token }
+        #         yield job_id, intervals, NetFlowBot.perform_job, job_params
+
+        # mock the jobs for now: (until frontend is done)
+        job_id = 'traffic_per_protocol'
+        intervals = [60]
+        job_params = {
+            "entity_info": {
+                "account_id": 129104112,
+                "entity_id": 123,
+            },
+            "backend_url": self.backend_url,
+            "bot_token": self.bot_token,
+        }
+        yield job_id, intervals, NetFlowBot.perform_job, job_params
+
+    # This method is called whenever the job needs to be done. It gets the parameters and performs fetching of data.
+    @staticmethod
+    def perform_job(*args, **job_params):
+        traffic_per_protocol = r.hgetall(f'{REDIS_PREFIX}_{REDIS_HASH_TRAFFIC_PER_PROTOCOL}')
+        entity_info = job_params["entity_info"]
+        values = []
+        now = time.time()
+        for protocol, traffic_counter in traffic_per_protocol.items():
+            output_path = f'entity.{entity_info["entity_id"]}.netflow.traffic_per_protocol.{protocol.decode("utf-8")}'
+
+            # since we are getting the counters, convert them to values:
+            dv, dt = convert_counter_to_value(f'{REDIS_PREFIX}_counter_{output_path}', traffic_counter, now)
+            if dv is None:
+                continue
+            values.append({
+                'p': output_path,
+                'v': dv / dt,
+            })
+
+        if not values:
+            log.warning("No values found to be sent to Grafolean")
+            return
+
+        # send the data to Grafolean:
+        send_results_to_grafolean(
+            job_params['backend_url'],
+            job_params['bot_token'],
+            job_params['entity_info']['account_id'],
+            values,
+        )
+
+
+def _get_previous_counter_value(counter_ident):
+    prev_value = r.hgetall(counter_ident)
+    if not prev_value:  # empty dict
+        return None, None
+    return int(prev_value[b'v']), float(prev_value[b't'])
+
+
+def _save_current_counter_value(new_value, now, counter_ident):
+    r.hmset(counter_ident, {b'v': new_value, b't': now})
+
+
+def convert_counter_to_value(counter_ident, new_value, now):
+    old_value, t = _get_previous_counter_value(counter_ident)
+    new_value = int(float(new_value))
+    _save_current_counter_value(new_value, now, counter_ident)
+    if old_value is None:
+        # no previous counter, can't calculate value:
+        log.debug(f"Counter {counter_ident} has no previous value.")
+        return None, None
+    if new_value < old_value:
+        # new counter value is lower than the old one, probably overflow: (or reset)
+        log.warning(f"Counter overflow detected for counter {counter_ident}, discarding value - if this happens often, consider decreasing polling interval.")
+        return None, None
+    dt = now - t
+    dv = new_value - old_value
+    return dv, dt
 
 
 def wait_for_grafolean(backend_url):

netflowcollector.py

@@ -16,7 +16,7 @@
 sys.path.append(os.path.dirname(os.path.realpath(__file__)) + '/pynetflow')
 from pynetflow.main import get_export_packets
 
-from lookup import PROTOCOLS
+from lookup import PROTOCOLS, REDIS_HASH_TRAFFIC_PER_PROTOCOL
 
 
 logging.basicConfig(format='%(asctime)s | %(levelname)s | %(message)s',
@@ -32,6 +32,9 @@
 r = redis.Redis(host=REDIS_HOST)
 
 
+REDIS_PREFIX = 'netflow'
+
+
 def process_netflow(netflow_port):
     for ts, client, export in get_export_packets('0.0.0.0', NETFLOW_PORT):
         data = defaultdict(int)
@@ -40,10 +43,10 @@ def process_netflow(netflow_port):
             in_bytes = flow.data['IN_BYTES']
 
             protocol_str = PROTOCOLS.get(protocol, f'?{protocol}')
-            data[f'protocol_{protocol_str}_traffic'] += in_bytes
+            data[protocol_str] += in_bytes
 
         for k, v in data.items():
-            r.incrby(k, v)
+            r.hincrby(f'{REDIS_PREFIX}_{REDIS_HASH_TRAFFIC_PER_PROTOCOL}', k, v)
 
 
 if __name__ == "__main__":