move zeek func to zeek.py

src/navv/commands.py

@@ -1,5 +1,4 @@
 """CLI Commands."""
-import json
 import os
 import webbrowser
 
@@ -25,8 +24,8 @@
     write_stats_sheet,
     write_unknown_internals_sheet,
 )
-from navv.zeek import run_zeek, perform_zeekcut
-from navv.utilities import pushd, trim_dns_data
+from navv.zeek import get_dns_data, run_zeek, perform_zeekcut
+from navv.utilities import pushd
 
 
 @click.command("generate")
@@ -76,16 +75,8 @@ def generate(customer_name, output_dir, pcap, zeek_logs):
     # Get dns data for resolution
     json_path = os.path.join(output_dir, f"{customer_name}_dns_data.json")
 
-    dns_filtered = {}
-    if os.path.exists(json_path):
-        with open(json_path, "rb") as json_file:
-            dns_filtered = json.load(json_file)
-    else:
-        dns_data = perform_zeekcut(
-            fields=["query", "answers", "qtype", "rcode_name"],
-            log_file=os.path.join(zeek_logs, "dns.log"),
-        )
-        dns_filtered = trim_dns_data(dns_data)
+    # Get dns data from zeek logs
+    dns_filtered = get_dns_data(customer_name, output_dir, zeek_logs)
 
     # Get zeek dataframe
     zeek_df = get_zeek_df(zeek_data, dns_filtered)

src/navv/zeek.py

@@ -1,7 +1,9 @@
+import json
+import os
 from subprocess import Popen, PIPE, STDOUT, check_call
 
 from navv.message_handler import error_msg
-from navv.utilities import pushd, timeit
+from navv.utilities import pushd, timeit, trim_dns_data
 
 
 @timeit
@@ -14,6 +16,21 @@ def run_zeek(pcap_path, zeek_logs_path, **kwargs):
             error_msg(e)
 
 
+@timeit
+def get_dns_data(customer_name, output_dir, zeek_logs):
+    """Get DNS data from zeek logs or from a json file if it exists"""
+    json_path = os.path.join(output_dir, f"{customer_name}_dns_data.json")
+    if os.path.exists(json_path):
+        with open(json_path, "rb") as json_file:
+            return json.load(json_file)
+
+    dns_data = perform_zeekcut(
+        fields=["query", "answers", "qtype", "rcode_name"],
+        log_file=os.path.join(zeek_logs, "dns.log"),
+    )
+    return trim_dns_data(dns_data)
+
+
 def perform_zeekcut(fields, log_file):
     """Perform the call to zeek-cut with the identified fields on the specified log file"""
     try: