Change handling of output dir for created trace XML and fix Label Bug

code/Attack/MembersMgmtCommAttack.py

@@ -4,6 +4,7 @@ from collections import deque
 from scipy.stats import gamma
 from lea import Lea
 from datetime import datetime
+import os
 
 from Attack import BaseAttack
 from Attack.AttackParameters import Parameter as Param
@@ -148,7 +149,7 @@ class MembersMgmtCommAttack(BaseAttack.BaseAttack):
         
     def generate_attack_pcap(self, context):
         # create the final messages that have to be sent, including all bot configurations
-        messages = self._create_messages()
+        messages = self._create_messages(context)
 
         if messages == []:
             return 0, []
@@ -225,7 +226,7 @@ class MembersMgmtCommAttack(BaseAttack.BaseAttack):
         return total_pkts , path_attack_pcap
 
 
-    def _create_messages(self):
+    def _create_messages(self, context):
         def add_ids_to_config(ids_to_add: list, existing_ips: list, new_ips: list, bot_configs: dict, idtype:str="local", router_mac:str=""):
             """
             Creates IP and MAC configurations for the given IDs and adds them to the existing configurations object.
@@ -313,6 +314,23 @@ class MembersMgmtCommAttack(BaseAttack.BaseAttack):
             return timestamp + minDelay + general_offset + unique_offset
 
 
+        def move_xml_to_outdir(filepath_xml: str):
+            """
+            Moves the XML file at filepath_xml to the output directory of the PCAP
+            :param filepath_xml: the filepath to the XML file
+            :return: the new filepath to the XML file
+            """
+
+            pcap_dir = context.get_output_dir()
+            xml_name = os.path.basename(filepath_xml)
+            if pcap_dir.endswith("/"):
+                new_xml_path = pcap_dir + xml_name
+            else:
+                new_xml_path = pcap_dir + "/" + xml_name
+            os.rename(filepath_xml, new_xml_path)
+            context.add_other_created_file(new_xml_path)
+            return new_xml_path
+
         # parse input CSV or XML
         filepath_xml = self.get_param_value(Param.FILE_XML)
         filepath_csv = self.get_param_value(Param.FILE_CSV)
@@ -320,6 +338,8 @@ class MembersMgmtCommAttack(BaseAttack.BaseAttack):
         # prefer XML input over CSV input (in case both are given)
         if filepath_csv and filepath_xml == self.DEFAULT_XML_PATH:
             filepath_xml = FileUtils.parse_csv_to_xml(filepath_csv) 
+            filepath_xml = move_xml_to_outdir(filepath_xml)
+
 
         abstract_packets = FileUtils.parse_xml(filepath_xml)
 

code/ID2TLib/AttackContext.py

@@ -1,11 +1,13 @@
 import tempfile
 
 class AttackContext:
-    def __init__(self):
+    def __init__(self, out_dir):
         # files allocated by attacks, intended to get copied next to the pcap
         # the keys are the suffix for later use, e.g. "_extra_info.txt"
         # the values are the respectable file names
         self.allocated_files = {}
+        self.other_created_files = set()
+        self.out_dir = out_dir
 
     def allocate_file(self, suffix):
         if suffix in self.allocated_files:
@@ -16,6 +18,12 @@ class AttackContext:
 
         return file
 
+    def add_other_created_file(self, filepath):
+        self.other_created_files.add(filepath)
+
+    def get_other_created_files(self):
+        return sorted(self.other_created_files)
+
     def reset(self):
         self.allocated_files.clear()
 
@@ -24,3 +32,6 @@ class AttackContext:
 
         self.allocated_files.clear()
         return return_
+
+    def get_output_dir(self):
+        return self.out_dir

code/ID2TLib/AttackController.py

@@ -4,7 +4,7 @@ import sys
 from Attack.AttackParameters import Parameter
 from ID2TLib import LabelManager
 from ID2TLib import Statistics
-from ID2TLib.Label import Label
+from ID2TLib.OldLibs.Label import Label
 from ID2TLib.PcapFile import PcapFile
 
 

code/ID2TLib/Controller.py

@@ -55,26 +55,24 @@ class Controller:
         :param attacks_config: A list of attacks with their attack parameters.
         """
 
+        # get output directory
+        if self.pcap_out_path:
+            out_dir = os.path.dirname(self.pcap_out_path)
+        else:
+            out_dir = os.path.dirname(self.pcap_src_path)
+        # if out_dir is cwd
+        if out_dir == "":
+            out_dir = "."
+
         # context for the attack(s)
-        context = AttackContext()
+        context = AttackContext(out_dir)
 
         # note if new xml file has been created by MembersMgmtCommAttack
-        created_xml = None
         # load attacks sequentially
         for attack in attacks_config:
-            # check if new xml file has been created by MembersMgmtCommAttack
-            if attack[0] == "MembersMgmtCommAttack":
-                for param in attack[1:]:
-                    key, value = param.split("=")
-                    if key == "file.csv":
-                        if os.path.isfile(value):
-                            created_xml, _ = os.path.splitext(value)
-                            created_xml += ".xml"
-                            break
             temp_attack_pcap = self.attack_controller.process_attack(attack[0], attack[1:], context)
             self.written_pcaps.append(temp_attack_pcap)
 
-
         # merge attack pcaps to get single attack pcap
         if len(self.written_pcaps) > 1:
             print("\nMerging temporary attack pcaps into single pcap file...", end=" ")
@@ -110,29 +108,24 @@ class Controller:
         # write label file with attacks
         self.label_manager.write_label_file(self.pcap_dest_path)
 
-        # if MembersMgmtCommAttack created an xml file, move it into input pcap directory
-        if created_xml:
-            pcap_dir = os.path.splitext(self.pcap_dest_path)[0]
-            if "/" in pcap_dir:
-                pcap_dir = "/".join(pcap_dir.split("/")[:-1])
-            xml_name = os.path.splitext(created_xml)[0] + ".xml"
-            if "/" in xml_name:
-                xml_name = xml_name.split("/")[-1]
-            new_xml_path = pcap_dir + "/" + xml_name
-            os.rename(created_xml, new_xml_path)
-
         # pcap_base contains the name of the pcap-file without the ".pcap" extension
         pcap_base = os.path.splitext(self.pcap_dest_path)[0]
+        created_files = [self.pcap_dest_path, self.label_manager.label_file_path]
         for suffix, filename in context.get_allocated_files():
             print(filename, pcap_base + suffix)
             shutil.move(filename, pcap_base + suffix)
+            created_files.append(pcap_base + suffix)
         context.reset()
 
         # print status message
-        if created_xml:
-            print('\nOutput files created: \n', self.pcap_dest_path, '\n', self.label_manager.label_file_path, '\n', new_xml_path)
-        else:
-            print('\nOutput files created: \n', self.pcap_dest_path, '\n', self.label_manager.label_file_path)
+        created_files += context.get_other_created_files()
+        created_files.sort()
+        print("\nOutput files created:")
+        for file in created_files:
+            # remove ./ at beginning of file to have only one representation for cwd
+            if file.startswith("./"):
+                file = file[2:]
+            print(file)
 
     def process_db_queries(self, query, print_results=False):
         """