Browse Source

Refactor creation of misc output files in MembersMgmtCommAttack

dustin.born 6 years ago
parent
commit
53748600d7
1 changed files with 11 additions and 28 deletions
  1. 11 28
      code/Attack/MembersMgmtCommAttack.py

+ 11 - 28
code/Attack/MembersMgmtCommAttack.py

@@ -12,6 +12,7 @@ from Attack import BaseAttack
 from Attack.AttackParameters import Parameter as Param
 from Attack.AttackParameters import Parameter as Param
 from Attack.AttackParameters import ParameterTypes
 from Attack.AttackParameters import ParameterTypes
 from ID2TLib.Ports import PortSelectors
 from ID2TLib.Ports import PortSelectors
+import ID2TLib.Utility as Util
 
 
 class MessageType(Enum):
 class MessageType(Enum):
     """
     """
@@ -176,15 +177,14 @@ class MembersMgmtCommAttack(BaseAttack.BaseAttack):
 
 
         self.add_param_value(Param.HIDDEN_MARK, False)
         self.add_param_value(Param.HIDDEN_MARK, False)
 
 
-    def generate_attack_pcap(self, context):
+    def generate_attack_pcap(self):
         """
         """
         Injects the packets of this attack into a PCAP and stores it as a temporary file.
         Injects the packets of this attack into a PCAP and stores it as a temporary file.
-        :param context: the context of the attack, containing e.g. files that are to be created
         :return: a tuple of the number packets injected and the path to the temporary attack PCAP
         :return: a tuple of the number packets injected and the path to the temporary attack PCAP
         """
         """
 
 
         # create the final messages that have to be sent, including all bot configurations
         # create the final messages that have to be sent, including all bot configurations
-        messages = self._create_messages(context)
+        messages = self._create_messages()
 
 
         if messages == []:
         if messages == []:
             return 0, []
             return 0, []
@@ -280,7 +280,10 @@ class MembersMgmtCommAttack(BaseAttack.BaseAttack):
                 last_packet = packets[-1]
                 last_packet = packets[-1]
 
 
         # write the mapping to a file
         # write the mapping to a file
-        msg_packet_mapping.write_to(context.allocate_file("_mapping.xml"))
+        current_ts = datetime.now().strftime("%Y%m%d-%H%M%S")
+        mapping_filename = "mapping_" + current_ts 
+        msg_packet_mapping.write_to_file(mapping_filename)
+        Util.MISC_OUT_FILES["mapping.xml"] = mapping_filename
 
 
         # Store timestamp of last packet
         # Store timestamp of last packet
         self.attack_end_utime = last_packet.time
         self.attack_end_utime = last_packet.time
@@ -293,10 +296,9 @@ class MembersMgmtCommAttack(BaseAttack.BaseAttack):
         pass
         pass
 
 
 
 
-    def _create_messages(self, context):
+    def _create_messages(self):
         """
         """
         Creates the messages that are to be injected into the PCAP.
         Creates the messages that are to be injected into the PCAP.
-        :param context: the context of the attack, containing e.g. files that are to be created
         :return: the final messages as a list
         :return: the final messages as a list
         """
         """
 
 
@@ -478,24 +480,6 @@ class MembersMgmtCommAttack(BaseAttack.BaseAttack):
                 else:
                 else:
                     bot_configs[bot_id]["TTL"] = total_ttl_prob_dict.random()
                     bot_configs[bot_id]["TTL"] = total_ttl_prob_dict.random()
 
 
-        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
         # parse input CSV or XML
         filepath_xml = self.get_param_value(Param.FILE_XML)
         filepath_xml = self.get_param_value(Param.FILE_XML)
         filepath_csv = self.get_param_value(Param.FILE_CSV)
         filepath_csv = self.get_param_value(Param.FILE_CSV)
@@ -507,7 +491,7 @@ class MembersMgmtCommAttack(BaseAttack.BaseAttack):
         # --> prefer XML input over CSV input (in case both are given)
         # --> prefer XML input over CSV input (in case both are given)
         print_updates = False
         print_updates = False
         if filepath_csv and filepath_xml == self.DEFAULT_XML_PATH:
         if filepath_csv and filepath_xml == self.DEFAULT_XML_PATH:
-            filename = os.path.splitext(filepath_csv)[0]
+            filepath_xml = OUT_DIR + os.path.splitext(os.path.basename(filepath_csv))[0]
             filesize = os.path.getsize(filepath_csv) / 2**20  # get filesize in MB
             filesize = os.path.getsize(filepath_csv) / 2**20  # get filesize in MB
             if filesize > 10:
             if filesize > 10:
                 print("\nParsing input CSV file...", end=" ")
                 print("\nParsing input CSV file...", end=" ")
@@ -518,8 +502,8 @@ class MembersMgmtCommAttack(BaseAttack.BaseAttack):
                 print("done.")
                 print("done.")
                 print("Writing corresponding XML file...", end=" ")
                 print("Writing corresponding XML file...", end=" ")
                 sys.stdout.flush()
                 sys.stdout.flush()
-            filepath_xml = cpp_comm_proc.write_xml(filename)
-            filepath_xml = move_xml_to_outdir(filepath_xml)
+            filepath_xml = cpp_comm_proc.write_xml(filepath_xml)
+            Util.MISC_OUT_FILES[filepath_xml] = None
             if print_updates: print("done.")
             if print_updates: print("done.")
         else:
         else:
             filesize = os.path.getsize(filepath_xml) / 2**20  # get filesize in MB
             filesize = os.path.getsize(filepath_xml) / 2**20  # get filesize in MB
@@ -530,7 +514,6 @@ class MembersMgmtCommAttack(BaseAttack.BaseAttack):
             cpp_comm_proc.parse_xml(filepath_xml)
             cpp_comm_proc.parse_xml(filepath_xml)
             if print_updates: print("done.")
             if print_updates: print("done.")
 
 
-
         # find a good communication mapping in the input file that matches the users parameters
         # find a good communication mapping in the input file that matches the users parameters
         nat = self.get_param_value(Param.NAT_PRESENT)
         nat = self.get_param_value(Param.NAT_PRESENT)
         comm_proc = CommunicationProcessor(self.msg_types, nat)
         comm_proc = CommunicationProcessor(self.msg_types, nat)