From 5446ea11221a033f1b32a7ac6f24d62ecdbed059 Mon Sep 17 00:00:00 2001
From: Tomasz Grzegowski <tomasz.grzegowski@nordsec.com>
Date: Mon, 23 Dec 2024 14:19:33 +0000
Subject: [PATCH] Refactor result file handling

---
 xray/analyze.py | 17 ++++++++++-------
 xray/paths.py   | 24 ++++++++++++++++++++++++
 xray/run.py     | 16 +++++++++-------
 3 files changed, 43 insertions(+), 14 deletions(-)
 create mode 100644 xray/paths.py

diff --git a/xray/analyze.py b/xray/analyze.py
index 5123146..252c0ef 100644
--- a/xray/analyze.py
+++ b/xray/analyze.py
@@ -6,10 +6,11 @@
 from scapy.all import PcapReader  # type: ignore
 from scapy.layers.inet import UDP  # type: ignore
 import mpl_ascii
+from paths import PathGenerator
 
 
-def analyze(csv_path, test_path, count, test_type, ascii, save_output):
-    Analyzer(csv_path, test_path, count, test_type, ascii, save_output)
+def analyze(paths, count, test_type, ascii, save_output):
+    Analyzer(paths, count, test_type, ascii, save_output)
 
 
 class CsvData:
@@ -63,10 +64,12 @@ def __init__(self, pcap_path, test_type):
 
 
 class Analyzer:
-    def __init__(self, csv_name, test_path, count, test_type, ascii_output, save_output):
+    def __init__(
+        self, paths: PathGenerator, count, test_type, ascii_output, save_output
+    ):
         self.count = count
-        self.csv_data = CsvData(csv_name)
-        self.pcap_data = PcapData(test_path + ".pcap", test_type)
+        self.csv_data = CsvData(paths.csv())
+        self.pcap_data = PcapData(paths.pcap(), test_type)
 
         graphs = [
             self.ordering_pie_chart,
@@ -86,7 +89,7 @@ def __init__(self, csv_name, test_path, count, test_type, ascii_output, save_out
             fn(ax[i, 1])
 
         if save_output:
-            plt.savefig(test_path + ".png")
+            plt.savefig(paths.png())
         plt.show()
 
         if ascii_output:
@@ -108,7 +111,7 @@ def __init__(self, csv_name, test_path, count, test_type, ascii_output, save_out
                 fn(ax[i])
 
             if save_output:
-                plt.savefig(test_path + ".txt")
+                plt.savefig(paths.txt())
             plt.show()
 
     def ordering_pie_chart(self, ax):
diff --git a/xray/paths.py b/xray/paths.py
new file mode 100644
index 0000000..b61b45d
--- /dev/null
+++ b/xray/paths.py
@@ -0,0 +1,24 @@
+class PathGenerator:
+
+    def __init__(self, wg_ifc_name: str, wg: str, test_type: str, count: int):
+        self.wg_ifc_name = wg_ifc_name
+        self.wg = wg
+        self.test_type = test_type
+        self.count = count
+
+    def base_path(self) -> str:
+        return f"results/{self.wg_ifc_name}_{self.wg.lower()}_{self.test_type}_{self.count}"
+
+    def csv(self) -> str:
+        return (
+            f"results/xray_metrics_{self.wg.lower()}_{self.test_type}_{self.count}.csv"
+        )
+
+    def pcap(self) -> str:
+        return f"{self.base_path()}.pcap"
+
+    def txt(self) -> str:
+        return f"{self.base_path()}.txt"
+
+    def png(self) -> str:
+        return f"{self.base_path()}.png"
diff --git a/xray/run.py b/xray/run.py
index d41ce46..24639c2 100755
--- a/xray/run.py
+++ b/xray/run.py
@@ -7,6 +7,7 @@
 from analyze import analyze
 from enum import Enum
 from pathlib import Path
+from paths import PathGenerator
 
 WG_IFC_NAME = "xraywg1"
 
@@ -126,17 +127,19 @@ def main():
     build_neptun = args.nobuild_neptun is False
     build_xray = args.nobuild_xray is False
 
+    file_paths = PathGenerator(WG_IFC_NAME, wg.name, test_type, count)
+
     Path("results/").mkdir(parents=True, exist_ok=True)
     try:
-        os.remove(get_csv_name(wg.name, test_type, count))
-        os.remove(get_test_path(wg.name, test_type, count) + ".pcap")
-        os.remove(get_test_path(wg.name, test_type, count) + ".png")
-        os.remove(get_test_path(wg.name, test_type, count) + ".txt")
+        os.remove(file_paths.csv())
+        os.remove(file_paths.pcap())
+        os.remove(file_paths.png())
+        os.remove(file_paths.txt())
     except:  # noqa: E722
         pass
 
     setup_wireguard(wg, build_neptun, args.disable_drop_privileges)
-    tcpdump = start_tcpdump(get_test_path(wg.name, test_type, count) + ".pcap")
+    tcpdump = start_tcpdump(file_paths.pcap())
 
     succeeded = True
     try:
@@ -150,8 +153,7 @@ def main():
 
     if succeeded:
         analyze(
-            get_csv_name(wg.name, test_type, count),
-            get_test_path(wg.name, test_type, count),
+            file_paths,
             count,
             test_type,
             args.ascii,