zeek.py: log the time each func takes inside process_line and the db calls in add_flow_to_profile

Author: AlyaGomaa

slips_files/core/input_profilers/zeek.py

@@ -1,5 +1,8 @@
 # SPDX-FileCopyrightText: 2021 Sebastian Garcia <[email protected]>
 # SPDX-License-Identifier: GPL-2.0-only
+import csv
+import os
+import time
 from re import split
 from typing import Dict
 from slips_files.common.abstracts.iinput_type import IInputType
@@ -148,6 +151,32 @@ def fill_empty_class_fields(self, flow_values: dict, slips_class):
 class ZeekJSON(IInputType, Zeek):
     def __init__(self):
         self.line_processor_cache = {}
+        self.times = {}
+        self.init_csv()
+
+    def init_csv(self):
+        path = os.path.join("/tmp/", "zeek_json_times_each_func_took.csv")
+        with open(path, "w", newline="") as f:
+            writer = csv.writer(f)
+            writer.writerow(
+                [
+                    "get_file_type",
+                    "removing_empty_vals",
+                    "setting_conn_vals_to_0",
+                    "fill_empty_class_fields",
+                    "calling_slips_class",
+                ]
+            )
+
+    def log_time(self, what, time):
+        self.times[what] = f"{time:.2f}"
+        last_metric = "calling_slips_class"
+        if what == last_metric:
+            path = os.path.join("/tmp/", "zeek_json_times_each_func_took.csv")
+            with open(path, "a", newline="") as f:
+                writer = csv.DictWriter(f, fieldnames=self.times.keys())
+                writer.writerow(self.times)
+            self.times = {}
 
     def process_line(self, new_line: dict):
         line = new_line["data"]
@@ -156,7 +185,11 @@ def process_line(self, new_line: dict):
         if not isinstance(line, dict):
             return False
 
+        n = time.time()
         file_type = self.get_file_type(new_line)
+        latency = time.time() - n
+        self.log_time("get_file_type", latency)
+
         line_map = LOG_MAP.get(file_type)
         if not line_map:
             return False
@@ -168,6 +201,7 @@ def process_line(self, new_line: dict):
 
         flow_values = {"starttime": starttime, "interface": interface}
 
+        n = time.time()
         for zeek_field, slips_field in line_map.items():
             if not slips_field:
                 continue
@@ -176,25 +210,37 @@ def process_line(self, new_line: dict):
                 val = ""
             flow_values[slips_field] = val
 
+        latency = time.time() - n
+        self.log_time("removing_empty_vals", latency)
+
         if file_type in LINE_TYPE_TO_SLIPS_CLASS:
+            n = time.time()
             slips_class = LINE_TYPE_TO_SLIPS_CLASS[file_type]
-
             if file_type == "conn.log":
                 flow_values["dur"] = float(flow_values.get("dur", 0) or 0)
-                for fld in (
+                for field in (
                     "sbytes",
                     "dbytes",
                     "spkts",
                     "dpkts",
                     "sport",
                     "dport",
                 ):
-                    flow_values[fld] = int(flow_values.get(fld, 0) or 0)
-
+                    flow_values[field] = int(flow_values.get(field, 0) or 0)
+                latency = time.time() - n
+                self.log_time("setting_conn_vals_to_0", latency)
+            n = time.time()
             flow_values = self.fill_empty_class_fields(
                 flow_values, slips_class
             )
+            latency = time.time() - n
+            self.log_time("fill_empty_class_fields", latency)
+
+            n = time.time()
             self.flow = slips_class(**flow_values)
+            latency = time.time() - n
+            self.log_time("calling_slips_class", latency)
+
             return self.flow
 
         print(f"[Profiler] Invalid file_type: {file_type}, line: {line}")

slips_files/core/profiler_worker.py

@@ -531,22 +531,30 @@ def add_flow_to_profile(self, flow):
         time_it_Took = time.time() - n
         self.log_time("is_whitelisted_flow", time_it_Took)
 
-        n = time.time()
-
         # 5th. Store the data according to the paremeters
         # Now that we have the profileid and twid, add the data from the flow
         # in this tw for this profile
         self.print(f"Storing data in the profile: {profileid}", 3, 0)
+
+        n = time.time()
         flow.starttime = self.convert_starttime_to_epoch(flow.starttime)
+        time_it_Took = time.time() - n
+
+        self.log_time("convert_starttime_to_epoch", time_it_Took)
         # For this 'forward' profile, find the id in the
-        # database of the tw where the flow belongs.
+        # database of the tw where the flow belongs.        n = time.time()
+        n = time.time()
         twid = self.db.get_timewindow(flow.starttime, profileid)
+        time_it_Took = time.time() - n
+        self.log_time("get_timewindow", time_it_Took)
+
         flow_parser.twid = twid
 
+        n = time.time()
         # Create profiles for all ips we see
         self.db.add_profile(profileid, flow.starttime)
         time_it_Took = time.time() - n
-        self.log_time("db_calls", time_it_Took)
+        self.log_time("add_profile", time_it_Took)
 
         n = time.time()
         self.store_features_going_out(flow, flow_parser)
@@ -570,7 +578,9 @@ def init_csv(self):
                 [
                     "get_gateway_info",
                     "is_whitelisted_flow",
-                    "db_calls",
+                    "convert_starttime_to_epoch",
+                    "get_timewindow",
+                    "add_profile",
                     "store_features_going_out",
                     "process_line",
                     "add_flow_to_profile",