First working version of label-mapping (now with 30% more bugs)

code/Attack/BaseAttack.py

@@ -453,6 +453,12 @@ class BaseAttack(metaclass=ABCMeta):
 
         return destination
 
+    def post_pcap_written(self, final_filename):
+        """
+        :param final_filename: The filename of the final pcap created
+        """
+        pass
+
     def get_reply_delay(self, ip_dst):
         """
            Gets the minimum and the maximum reply delay for all the connections of a specific IP.

code/Attack/DDoSAttack.py

@@ -71,7 +71,7 @@ class DDoSAttack(BaseAttack.BaseAttack):
         self.add_param_value(Param.MAC_DESTINATION, destination_mac)
         self.add_param_value(Param.VICTIM_BUFFER, randint(1000,10000))
 
-    def generate_attack_pcap(self):
+    def generate_attack_pcap(self, context):
         def update_timestamp(timestamp, pps, delay=0):
             """
             Calculates the next timestamp to be used based on the packet per second rate (pps) and the maximum delay.

code/Attack/EternalBlueExploit.py

@@ -77,7 +77,7 @@ class EternalBlueExploit(BaseAttack.BaseAttack):
                               self.statistics.get_pps_received(most_used_ip_address)) / 2)
         self.add_param_value(Param.INJECT_AFTER_PACKET, randint(0, self.statistics.get_packet_count()))
 
-    def generate_attack_pcap(self):
+    def generate_attack_pcap(self, context):
         def update_timestamp(timestamp, pps):
             """
             Calculates the next timestamp to be used based on the packet per second rate (pps) and the maximum delay.

code/Attack/JoomlaRegPrivExploit.py

@@ -78,7 +78,7 @@ class JoomlaRegPrivExploit(BaseAttack.BaseAttack):
                              (self.statistics.get_pps_sent(most_used_ip_address) +
                               self.statistics.get_pps_received(most_used_ip_address)) / 2)
 
-    def generate_attack_pcap(self):
+    def generate_attack_pcap(self, context):
         def update_timestamp(timestamp, pps):
             """
             Calculates the next timestamp to be used based on the packet per second rate (pps) and the maximum delay.

code/Attack/MembersMgmtCommAttack.py

@@ -12,10 +12,12 @@ class MessageType(Enum):
     SALITY_HELLO_REPLY = 104
 
 class Message():
+    INVALID_LINENO = -1
+
     """
     Defines a compact message type that contains all necessary information.
     """
-    def __init__(self, msg_id: int, src, dst, type_: MessageType, time: float, refer_msg_id: int=-1):
+    def __init__(self, msg_id: int, src, dst, type_: MessageType, time: float, refer_msg_id: int=-1, line_no = -1):
         """
         Constructs a message with the given parameters.
 
@@ -25,6 +27,7 @@ class Message():
         :param type_: the type of the message
         :param time: the timestamp of the message
         :param refer_msg_id: the ID this message is a request for or reply to. -1 if there is no related message.
+        :param line_no: The line number this message appeared in the original file
         """
         self.msg_id = msg_id
         self.src = src
@@ -32,6 +35,8 @@ class Message():
         self.type = type_
         self.time = time
         self.refer_msg_id = refer_msg_id
+        # if similar fields to line_no should be added consider a separate class
+        self.line_no = line_no
 
     def __str__(self):
         str_ = "{0}. at {1}: {2}-->{3}, {4}, refer:{5}".format(self.msg_id, self.time, self.src, self.dst, self.type, self.refer_msg_id)
@@ -54,6 +59,7 @@ from ID2TLib.PcapAddressOperations import PcapAddressOperations
 from ID2TLib.CommunicationProcessor import CommunicationProcessor
 from ID2TLib.MacAddressGenerator import MacAddressGenerator
 from ID2TLib.PortGenerator import gen_random_server_port
+from ID2TLib.Botnet.MessageMapping import MessageMapping
 
 
 class MembersMgmtCommAttack(BaseAttack.BaseAttack):
@@ -141,7 +147,7 @@ class MembersMgmtCommAttack(BaseAttack.BaseAttack):
 
 
         
-    def generate_attack_pcap(self):
+    def generate_attack_pcap(self, context):
         # create the final messages that have to be sent, including all bot configurations
         messages = self._create_messages()
 
@@ -160,6 +166,10 @@ class MembersMgmtCommAttack(BaseAttack.BaseAttack):
         limit_duration = self.get_param_value(Param.ATTACK_DURATION)
         duration = 0
         path_attack_pcap = None
+
+        msg_packet_mapping = MessageMapping(messages)
+
+
         # create packets to write to PCAP file
         for msg in messages:
             # retrieve the source and destination configurations
@@ -193,6 +203,7 @@ class MembersMgmtCommAttack(BaseAttack.BaseAttack):
 
             packet.time = pcap_timestamp
             packets.append(packet)
+            msg_packet_mapping.map_message(msg, packet)
             total_pkts += 1
 
             # Store timestamp of first packet (for attack label)
@@ -212,6 +223,9 @@ class MembersMgmtCommAttack(BaseAttack.BaseAttack):
             path_attack_pcap = self.write_attack_pcap(packets, True, path_attack_pcap)
             last_packet = packets[-1]
 
+        # write the mapping to a file
+        msg_packet_mapping.write_to(context.allocate_file("_mapping.xml"))
+
         # Store timestamp of last packet
         self.attack_end_utime = last_packet.time
 

code/Attack/PortscanAttack.py

@@ -107,7 +107,7 @@ class PortscanAttack(BaseAttack.BaseAttack):
             port_dst_shuffled = ports_dst
         return port_dst_shuffled
 
-    def generate_attack_pcap(self):
+    def generate_attack_pcap(self, context):
         def update_timestamp(timestamp, pps, delay=0):
             """
             Calculates the next timestamp to be used based on the packet per second rate (pps) and the maximum delay.

code/Attack/SQLiAttack.py

@@ -78,7 +78,7 @@ class SQLiAttack(BaseAttack.BaseAttack):
                              (self.statistics.get_pps_sent(most_used_ip_address) +
                               self.statistics.get_pps_received(most_used_ip_address)) / 2)
 
-    def generate_attack_pcap(self):
+    def generate_attack_pcap(self, context):
         def update_timestamp(timestamp, pps):
             """
             Calculates the next timestamp to be used based on the packet per second rate (pps) and the maximum delay.

code/Attack/SalityBotnet.py

@@ -56,7 +56,7 @@ class SalityBotnet(BaseAttack.BaseAttack):
                              (self.statistics.get_pps_sent(most_used_ip_address) +
                               self.statistics.get_pps_received(most_used_ip_address)) / 2)
 
-    def generate_attack_pcap(self):
+    def generate_attack_pcap(self, context):
         def update_timestamp(timestamp, pps):
             """
             Calculates the next timestamp to be used based on the packet per second rate (pps) and the maximum delay.

code/ID2TLib/AttackContext.py

@@ -0,0 +1,26 @@
+import tempfile
+
+class AttackContext:
+    def __init__(self):
+        # 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 = {}
+
+    def allocate_file(self, suffix):
+        if suffix in self.allocated_files:
+            raise ValueError("File with suffix %s is already allocated" % suffix)
+
+        file = tempfile.NamedTemporaryFile(mode = "w", suffix = suffix, delete = False)
+        self.allocated_files[suffix] = file.name
+
+        return file
+
+    def reset(self):
+        self.allocated_files.clear()
+
+    def get_allocated_files(self):
+        return_ = list(self.allocated_files.items())
+
+        self.allocated_files.clear()
+        return return_

code/ID2TLib/AttackController.py

@@ -40,7 +40,7 @@ class AttackController:
         # Record the attack
         self.added_attacks.append(self.current_attack)
 
-    def process_attack(self, attack: str, params: str):
+    def process_attack(self, attack: str, params: str, context):
         """
         Takes as input the name of an attack (classname) and the attack parameters as string. Parses the string of
         attack parameters, creates the attack by writing the attack packets and returns the path of the written pcap.
@@ -81,7 +81,7 @@ class AttackController:
         # Write attack into pcap file
         print("Generating attack packets...", end=" ")
         sys.stdout.flush()  # force python to print text immediately
-        total_packets, temp_attack_pcap_path = self.current_attack.generate_attack_pcap()
+        total_packets, temp_attack_pcap_path = self.current_attack.generate_attack_pcap(context)
         print("done. (total: " + str(total_packets) + " pkts.)")
 
         # Store label into LabelManager

code/ID2TLib/Botnet/MessageMapping.py

@@ -0,0 +1,52 @@
+import os.path
+from xml.dom.minidom import *
+
+
+class MessageMapping:
+    TAG_MAPPING_GROUP = "mappings"
+    TAG_MAPPING = "mapping"
+
+    ATTR_ID = "id"
+    ATTR_LINENO = "line_number"
+    ATTR_HAS_PACKET = "mapped"
+
+    ATTR_PACKET_TIME = "packet_time"
+
+    def __init__(self, messages):
+        self.messages = messages
+        self.id_to_packet = {}
+
+    def map_message(self, message, packet):
+        self.id_to_packet[message.msg_id] = packet
+
+    def to_xml(self):
+        doc = Document()
+
+        mappings = doc.createElement(self.TAG_MAPPING_GROUP)
+        doc.appendChild(mappings)
+
+        for message in self.messages:
+            mapping = doc.createElement(self.TAG_MAPPING)
+            mapping.setAttribute(self.ATTR_ID, str(message.msg_id))
+            mapping.setAttribute(self.ATTR_LINENO, str(message.line_no))
+
+            packet = self.id_to_packet.get(message.msg_id)
+            mapping.setAttribute(self.ATTR_HAS_PACKET, "true" if packet is not None else "false")
+            if packet:
+                mapping.setAttribute(self.ATTR_PACKET_TIME, str(packet.time))
+
+            mappings.appendChild(mapping)
+
+        return doc
+
+    def write_to(self, buffer, close = True):
+        buffer.write(self.to_xml().toprettyxml())
+        if close: buffer.close()
+
+    def write_to_file(self, filename: str, *args, **kwargs):
+        self.write_to(open(filename, "w", *args, **kwargs))
+
+    def write_next_to_pcap_file(self, pcap_filename : str, mapping_ext = "_mapping.xml", *args, **kwargs):
+        pcap_base = os.path.splitext(pcap_filename)[0]
+
+        self.write_to_file(pcap_base + mapping_ext, *args, **kwargs)

code/ID2TLib/CommunicationProcessor.py

@@ -188,6 +188,7 @@ class CommunicationProcessor():
         # process every packet individually 
         for packet in self.packets:
             id_src, id_dst, msg_type, time = packet["Src"], packet["Dst"], int(packet["Type"]), float(packet["Time"])
+            lineno = packet.get("LineNumber", -1)
             # if if either one of the IDs is not mapped, continue
             if (id_src not in local_init_ids) and (id_dst not in local_init_ids):
                 continue
@@ -206,7 +207,7 @@ class CommunicationProcessor():
                     respnd_ids.add(id_dst)
                 # convert the abstract message into a message object to handle it better
                 msg_str = "{0}-{1}".format(id_src, id_dst)
-                msg = Message(msg_id, id_src, id_dst, msg_type, time)
+                msg = Message(msg_id, id_src, id_dst, msg_type, time, line_no = lineno)
                 msgs.append(msg)
                 prev_reqs[msg_str] = msg_id
                 msg_id += 1
@@ -226,7 +227,7 @@ class CommunicationProcessor():
                 # find the request message ID for this response and set its reference index
                 refer_idx = prev_reqs[msg_str]
                 msgs[refer_idx].refer_msg_id = msg_id
-                msg = Message(msg_id, id_src, id_dst, msg_type, time, refer_idx)
+                msg = Message(msg_id, id_src, id_dst, msg_type, time, refer_idx, lineno)
                 msgs.append(msg)
                 # remove the request to this response from storage
                 del(prev_reqs[msg_str])
@@ -240,9 +241,9 @@ class CommunicationProcessor():
                 if refer_idx is not None:
                     msgs[refer_idx].refer_msg_id = msg_id
                     if msgs[refer_idx].type == MessageType.SALITY_NL_REQUEST:
-                        msg = Message(msg_id, id_src, id_dst, MessageType.SALITY_NL_REPLY, time, refer_idx)
+                        msg = Message(msg_id, id_src, id_dst, MessageType.SALITY_NL_REPLY, time, refer_idx, lineno)
                     else:
-                        msg = Message(msg_id, id_src, id_dst, MessageType.SALITY_HELLO_REPLY, time, refer_idx)
+                        msg = Message(msg_id, id_src, id_dst, MessageType.SALITY_HELLO_REPLY, time, refer_idx, lineno)
                     msgs.append(msg)
                     # remove the request to this response from storage
                     del(prev_reqs[msg_str])

code/ID2TLib/Controller.py

@@ -1,10 +1,12 @@
 import os
 import sys
+import shutil
 
 from ID2TLib.AttackController import AttackController
 from ID2TLib.LabelManager import LabelManager
 from ID2TLib.PcapFile import PcapFile
 from ID2TLib.Statistics import Statistics
+from ID2TLib.AttackContext import AttackContext
 
 
 class Controller:
@@ -52,6 +54,10 @@ class Controller:
         input dataset.
         :param attacks_config: A list of attacks with their attack parameters.
         """
+
+        # context for the attack(s)
+        context = AttackContext()
+
         # note if new xml file has been created by MembersMgmtCommAttack
         created_xml = None
         # load attacks sequentially
@@ -65,7 +71,7 @@ class Controller:
                             created_xml, _ = os.path.splitext(value)
                             created_xml += ".xml"
                             break
-            temp_attack_pcap = self.attack_controller.process_attack(attack[0], attack[1:])
+            temp_attack_pcap = self.attack_controller.process_attack(attack[0], attack[1:], context)
             self.written_pcaps.append(temp_attack_pcap)
 
 
@@ -111,6 +117,13 @@ class Controller:
             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]
+        for suffix, filename in context.get_allocated_files():
+            print(filename, pcap_base + suffix)
+            shutil.move(filename, 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)

code/ID2TLib/FileUtils.py

@@ -8,7 +8,7 @@ def parse_xml(filepath: str):
 	It is assumed, that packets are placed on the second hierarchical level and packetinformation is encoded as attributes
 	
 	:param filepath: the path to the XML file to be parsed
- 	:return: a List of Dictionaries, each Dictionary contains the information of one packet
+	:return: a List of Dictionaries, each Dictionary contains the information of one packet
 	'''
 
 	tree = ElementTree.parse(filepath)
@@ -26,7 +26,7 @@ def parse_csv_to_xml(filepath: str):
 	Converts a CSV file into an XML file. Every entry is converted to a child with respective attributes of the root node
 	
 	:param filepath: the path to the CSV file to be parsed
- 	:return: a path to the newly created XML file 
+	:return: a path to the newly created XML file
 	'''
 
 	filename = os.path.splitext(filepath)[0]
@@ -36,14 +36,20 @@ def parse_csv_to_xml(filepath: str):
 
 	# parse the csvFile into reader
 	with open(filepath, "rt") as csvFile:
-	    reader = csv.reader(csvFile, delimiter=",")
+		reader = csv.reader(csvFile, delimiter=",")
 		# loop through the parsed file, creating packet-elements with the structure of the csvFile as attributes
-	    for line in reader:
-	        packet = ElementTree.SubElement(root, "packet")
-	        for element in line:
-	        	element = element.replace(" ", "")
-	        	key, value = element.split(":")
-		        packet.attrib[key] = value
+		lineno = -1 # lines start at zero
+		for line in reader:
+			lineno += 1
+			if not line:
+				continue
+
+			packet = ElementTree.SubElement(root, "packet")
+			for element in line:
+				element = element.replace(" ", "")
+				key, value = element.split(":")
+				packet.attrib[key] = value
+			packet.attrib["LineNumber"] = lineno
 
 	# writing the ElementTree into the .xml file
 	tree = ElementTree.ElementTree(root)