refactor imports in Attack/EternalBlueExploit.py

refactor whitspaces in Attack/EternalBlueExploit.py
deep refactoring Attack/EternalBlueExploit.py

code/Attack/EternalBlueExploit.py

@@ -1,17 +1,17 @@
 import logging
-from random import randint, uniform
+import random as rnd
 
-from lea import Lea
-from scapy.layers.inet import Ether
-from scapy.utils import RawPcapReader
+import lea
+import scapy.layers.inet as inet
+import scapy.utils
 
-from Attack import BaseAttack
-from Attack.AttackParameters import Parameter as Param
-from Attack.AttackParameters import ParameterTypes
-from ID2TLib.SMBLib import smb_port
+import Attack.AttackParameters as atkParam
+import Attack.BaseAttack as BaseAttack
+import ID2TLib.SMBLib as SMBLib
 import ID2TLib.Utility as Util
 
 logging.getLogger("scapy.runtime").setLevel(logging.ERROR)
+
 # noinspection PyPep8
 
 
@@ -26,23 +26,25 @@ class EternalBlueExploit(BaseAttack.BaseAttack):
     def __init__(self):
         """
         Creates a new instance of the EternalBlue Exploit.
-
         """
         # Initialize attack
         super(EternalBlueExploit, self).__init__("EternalBlue Exploit", "Injects an EternalBlue exploit'",
-                                        "Privilege elevation")
+                                                 "Privilege elevation")
+
+        self.pkt_num = 0
+        self.path_attack_pcap = None
 
         # Define allowed parameters and their type
         self.supported_params.update({
-            Param.MAC_SOURCE: ParameterTypes.TYPE_MAC_ADDRESS,
-            Param.IP_SOURCE: ParameterTypes.TYPE_IP_ADDRESS,
-            Param.PORT_SOURCE: ParameterTypes.TYPE_PORT,
-            Param.MAC_DESTINATION: ParameterTypes.TYPE_MAC_ADDRESS,
-            Param.IP_DESTINATION: ParameterTypes.TYPE_IP_ADDRESS,
-            Param.PORT_DESTINATION: ParameterTypes.TYPE_PORT,
-            Param.INJECT_AT_TIMESTAMP: ParameterTypes.TYPE_FLOAT,
-            Param.INJECT_AFTER_PACKET: ParameterTypes.TYPE_PACKET_POSITION,
-            Param.PACKETS_PER_SECOND: ParameterTypes.TYPE_FLOAT
+            atkParam.Parameter.MAC_SOURCE: atkParam.ParameterTypes.TYPE_MAC_ADDRESS,
+            atkParam.Parameter.IP_SOURCE: atkParam.ParameterTypes.TYPE_IP_ADDRESS,
+            atkParam.Parameter.PORT_SOURCE: atkParam.ParameterTypes.TYPE_PORT,
+            atkParam.Parameter.MAC_DESTINATION: atkParam.ParameterTypes.TYPE_MAC_ADDRESS,
+            atkParam.Parameter.IP_DESTINATION: atkParam.ParameterTypes.TYPE_IP_ADDRESS,
+            atkParam.Parameter.PORT_DESTINATION: atkParam.ParameterTypes.TYPE_PORT,
+            atkParam.Parameter.INJECT_AT_TIMESTAMP: atkParam.ParameterTypes.TYPE_FLOAT,
+            atkParam.Parameter.INJECT_AFTER_PACKET: atkParam.ParameterTypes.TYPE_PACKET_POSITION,
+            atkParam.Parameter.PACKETS_PER_SECOND: atkParam.ParameterTypes.TYPE_FLOAT
         })
 
     def init_params(self):
@@ -50,8 +52,6 @@ class EternalBlueExploit(BaseAttack.BaseAttack):
         Initialize the parameters of this attack using the user supplied command line parameters.
         Use the provided statistics to calculate default parameters and to process user
         supplied queries.
-
-        :param statistics: Reference to a statistics object.
         """
         # PARAMETERS: initialize with default utilsvalues
         # (values are overwritten if user specifies them)
@@ -60,76 +60,76 @@ class EternalBlueExploit(BaseAttack.BaseAttack):
         random_ip_address = self.statistics.get_random_ip_address()
         while random_ip_address == most_used_ip_address:
             random_ip_address = self.statistics.get_random_ip_address()
-        self.add_param_value(Param.IP_SOURCE, random_ip_address)
-        self.add_param_value(Param.MAC_SOURCE, self.statistics.get_mac_address(random_ip_address))
-        self.add_param_value(Param.PORT_SOURCE, randint(self.minDefaultPort, self.maxDefaultPort))
+        self.add_param_value(atkParam.Parameter.IP_SOURCE, random_ip_address)
+        self.add_param_value(atkParam.Parameter.MAC_SOURCE, self.statistics.get_mac_address(random_ip_address))
+        self.add_param_value(atkParam.Parameter.PORT_SOURCE, rnd.randint(self.minDefaultPort, self.maxDefaultPort))
 
         # Victim configuration
-        self.add_param_value(Param.IP_DESTINATION, most_used_ip_address)
+        self.add_param_value(atkParam.Parameter.IP_DESTINATION, most_used_ip_address)
         destination_mac = self.statistics.get_mac_address(most_used_ip_address)
         if isinstance(destination_mac, list) and len(destination_mac) == 0:
             destination_mac = self.generate_random_mac_address()
-        self.add_param_value(Param.MAC_DESTINATION, destination_mac)
-        self.add_param_value(Param.PORT_DESTINATION, smb_port)
+        self.add_param_value(atkParam.Parameter.MAC_DESTINATION, destination_mac)
+        self.add_param_value(atkParam.Parameter.PORT_DESTINATION, SMBLib.smb_port)
 
         # Attack configuration
-        self.add_param_value(Param.PACKETS_PER_SECOND,
+        self.add_param_value(atkParam.Parameter.PACKETS_PER_SECOND,
                              (self.statistics.get_pps_sent(most_used_ip_address) +
                               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()))
+        self.add_param_value(atkParam.Parameter.INJECT_AFTER_PACKET, rnd.randint(0, self.statistics.get_packet_count()))
 
     def generate_attack_packets(self):
 
         # Timestamp
-        timestamp_next_pkt = self.get_param_value(Param.INJECT_AT_TIMESTAMP)
-        pps = self.get_param_value(Param.PACKETS_PER_SECOND)
+        timestamp_next_pkt = self.get_param_value(atkParam.Parameter.INJECT_AT_TIMESTAMP)
+        pps = self.get_param_value(atkParam.Parameter.PACKETS_PER_SECOND)
 
         # calculate complement packet rates of BG traffic per interval
         complement_interval_pps = self.statistics.calculate_complement_packet_rates(pps)
 
         # Initialize parameters
         self.packets = []
-        mac_source = self.get_param_value(Param.MAC_SOURCE)
-        ip_source = self.get_param_value(Param.IP_SOURCE)
-        port_source = self.get_param_value(Param.PORT_SOURCE)
-        mac_destination = self.get_param_value(Param.MAC_DESTINATION)
-        ip_destination = self.get_param_value(Param.IP_DESTINATION)
-        port_destination = self.get_param_value(Param.PORT_DESTINATION)
+        mac_source = self.get_param_value(atkParam.Parameter.MAC_SOURCE)
+        ip_source = self.get_param_value(atkParam.Parameter.IP_SOURCE)
+        # FIXME: why is port_source never used?
+        port_source = self.get_param_value(atkParam.Parameter.PORT_SOURCE)
+        mac_destination = self.get_param_value(atkParam.Parameter.MAC_DESTINATION)
+        ip_destination = self.get_param_value(atkParam.Parameter.IP_DESTINATION)
+        port_destination = self.get_param_value(atkParam.Parameter.PORT_DESTINATION)
 
         # Check ip.src == ip.dst
         self.ip_src_dst_equal_check(ip_source, ip_destination)
 
-        self.path_attack_pcap = None
-
         # Set TTL based on TTL distribution of IP address
         source_ttl_dist = self.statistics.get_ttl_distribution(ip_source)
         if len(source_ttl_dist) > 0:
-            source_ttl_prob_dict = Lea.fromValFreqsDict(source_ttl_dist)
+            source_ttl_prob_dict = lea.Lea.fromValFreqsDict(source_ttl_dist)
             source_ttl_value = source_ttl_prob_dict.random()
         else:
             source_ttl_value = Util.handle_most_used_outputs(self.statistics.process_db_query("most_used(ttlValue)"))
 
         destination_ttl_dist = self.statistics.get_ttl_distribution(ip_destination)
         if len(destination_ttl_dist) > 0:
-            destination_ttl_prob_dict = Lea.fromValFreqsDict(destination_ttl_dist)
+            destination_ttl_prob_dict = lea.Lea.fromValFreqsDict(destination_ttl_dist)
             destination_ttl_value = destination_ttl_prob_dict.random()
         else:
-            destination_ttl_value = Util.handle_most_used_outputs(self.statistics.process_db_query("most_used(ttlValue)"))
+            destination_ttl_value = Util.handle_most_used_outputs(
+                self.statistics.process_db_query("most_used(ttlValue)"))
 
         # Set Window Size based on Window Size distribution of IP address
         source_win_dist = self.statistics.get_win_distribution(ip_source)
         if len(source_win_dist) > 0:
-            source_win_prob_dict = Lea.fromValFreqsDict(source_win_dist)           
+            source_win_prob_dict = lea.Lea.fromValFreqsDict(source_win_dist)
         else:
-            source_win_dist =  self.statistics.get_win_distribution(self.statistics.get_most_used_ip_address())
-            source_win_prob_dict = Lea.fromValFreqsDict(source_win_dist)
-      
+            source_win_dist = self.statistics.get_win_distribution(self.statistics.get_most_used_ip_address())
+            source_win_prob_dict = lea.Lea.fromValFreqsDict(source_win_dist)
+
         destination_win_dist = self.statistics.get_win_distribution(ip_destination)
         if len(destination_win_dist) > 0:
-            destination_win_prob_dict = Lea.fromValFreqsDict(destination_win_dist)
+            destination_win_prob_dict = lea.Lea.fromValFreqsDict(destination_win_dist)
         else:
             destination_win_dist = self.statistics.get_win_distribution(self.statistics.get_most_used_ip_address())
-            destination_win_prob_dict = Lea.fromValFreqsDict(destination_win_dist)
+            destination_win_prob_dict = lea.Lea.fromValFreqsDict(destination_win_dist)
 
         # Set MSS (Maximum Segment Size) based on MSS distribution of IP address
         mss_value = Util.handle_most_used_outputs(self.statistics.process_db_query("most_used(mssValue)"))
@@ -139,37 +139,38 @@ class EternalBlueExploit(BaseAttack.BaseAttack):
         # Inject EternalBlue exploit packets
         # Read Win7_eternalblue_exploit pcap file
         source_origin_wins, destination_origin_wins = {}, {}
-        exploit_raw_packets = RawPcapReader(self.template_attack_pcap_path)
+        exploit_raw_packets = scapy.utils.RawPcapReader(self.template_attack_pcap_path)
 
-        port_source = randint(self.minDefaultPort,self.maxDefaultPort) # experiments show this range of ports
+        port_source = rnd.randint(self.minDefaultPort, self.maxDefaultPort)  # experiments show this range of ports
         # conversations = {(ip.src, ip.dst, port.src, port.dst): packets}
-        conversations, orderList_conversations = self.packets_to_convs(exploit_raw_packets)
+        conversations, order_list_conversations = self.packets_to_convs(exploit_raw_packets)
         exploit_raw_packets.close()
 
         conv_start_timesamp = timestamp_next_pkt
-        for conv_index, conv in enumerate(orderList_conversations):
-            conv_start_timesamp = conv_start_timesamp + uniform(0.001,0.01) # the distance between the starts of the converstaions
+        for conv_index, conv in enumerate(order_list_conversations):
+            # the distance between the starts of the converstaions
+            conv_start_timesamp = conv_start_timesamp + rnd.uniform(0.001, 0.01)
             timestamp_next_pkt = conv_start_timesamp
 
             conv_pkts = conversations[conv]
             inter_arrival_times = self.get_inter_arrival_time(conv_pkts)
 
-            if conv_index == len(orderList_conversations) - 2:  # Not the last conversation
-                timestamp_next_pkt = self.packets[-1].time + uniform(0.001,0.01)
+            if conv_index == len(order_list_conversations) - 2:  # Not the last conversation
+                timestamp_next_pkt = self.packets[-1].time + rnd.uniform(0.001, 0.01)
 
-            if conv_index != len(orderList_conversations)-1: # Not the last conversation
+            if conv_index != len(order_list_conversations) - 1:  # Not the last conversation
                 port_source += 2
                 for self.pkt_num, pkt in enumerate(conv_pkts):
-                    eth_frame = Ether(pkt[0])
+                    eth_frame = inet.Ether(pkt[0])
                     ip_pkt = eth_frame.payload
                     tcp_pkt = ip_pkt.payload
 
                     if self.pkt_num == 0:
-                        if tcp_pkt.getfieldval("dport") == smb_port:
+                        if tcp_pkt.getfieldval("dport") == SMBLib.smb_port:
                             orig_ip_dst = ip_pkt.getfieldval("dst")
 
                     # Request
-                    if ip_pkt.getfieldval("dst") == orig_ip_dst: # victim IP
+                    if ip_pkt.getfieldval("dst") == orig_ip_dst:  # victim IP
                         # Ether
                         eth_frame.setfieldval("src", mac_source)
                         eth_frame.setfieldval("dst", mac_destination)
@@ -180,13 +181,13 @@ class EternalBlueExploit(BaseAttack.BaseAttack):
                         # TCP
                         tcp_pkt.setfieldval("sport", port_source)
                         tcp_pkt.setfieldval("dport", port_destination)
-                        ## Window Size
+                        # Window Size
                         source_origin_win = tcp_pkt.getfieldval("window")
                         if source_origin_win not in source_origin_wins:
                             source_origin_wins[source_origin_win] = source_win_prob_dict.random()
                         new_win = source_origin_wins[source_origin_win]
                         tcp_pkt.setfieldval("window", new_win)
-                        ## MSS
+                        # MSS
                         tcp_options = tcp_pkt.getfieldval("options")
                         if tcp_options:
                             if tcp_options[0][0] == "MSS":
@@ -197,7 +198,8 @@ class EternalBlueExploit(BaseAttack.BaseAttack):
                         new_pkt.time = timestamp_next_pkt
 
                         pps = max(Util.get_interval_pps(complement_interval_pps, timestamp_next_pkt), 10)
-                        timestamp_next_pkt = Util.update_timestamp(timestamp_next_pkt, pps) + inter_arrival_times[self.pkt_num] #float(timeSteps.random())
+                        timestamp_next_pkt = Util.update_timestamp(timestamp_next_pkt, pps) + inter_arrival_times[
+                            self.pkt_num]  # float(timeSteps.random())
 
                     # Reply
                     else:
@@ -211,13 +213,13 @@ class EternalBlueExploit(BaseAttack.BaseAttack):
                         # TCP
                         tcp_pkt.setfieldval("dport", port_source)
                         tcp_pkt.setfieldval("sport", port_destination)
-                        ## Window Size
+                        # Window Size
                         destination_origin_win = tcp_pkt.getfieldval("window")
                         if destination_origin_win not in destination_origin_wins:
                             destination_origin_wins[destination_origin_win] = destination_win_prob_dict.random()
                         new_win = destination_origin_wins[destination_origin_win]
                         tcp_pkt.setfieldval("window", new_win)
-                        ## MSS
+                        # MSS
                         tcp_options = tcp_pkt.getfieldval("options")
                         if tcp_options:
                             if tcp_options[0][0] == "MSS":
@@ -227,17 +229,18 @@ class EternalBlueExploit(BaseAttack.BaseAttack):
                         new_pkt = (eth_frame / ip_pkt / tcp_pkt)
 
                         pps = max(Util.get_interval_pps(complement_interval_pps, timestamp_next_pkt), 10)
-                        timestamp_next_pkt = Util.update_timestamp(timestamp_next_pkt, pps) + inter_arrival_times[self.pkt_num]#float(timeSteps.random())
+                        timestamp_next_pkt = Util.update_timestamp(timestamp_next_pkt, pps) + inter_arrival_times[
+                            self.pkt_num]  # float(timeSteps.random())
 
                         new_pkt.time = timestamp_next_pkt
 
                     self.packets.append(new_pkt)
 
-            else: # Last conversation where the victim start a connection with the attacker
-                timestamp_next_pkt = self.packets[-1].time + uniform(0.001, 0.01)
-                port_source = randint(self.minDefaultPort,self.maxDefaultPort)
+            else:  # Last conversation where the victim start a connection with the attacker
+                timestamp_next_pkt = self.packets[-1].time + rnd.uniform(0.001, 0.01)
+                port_source = rnd.randint(self.minDefaultPort, self.maxDefaultPort)
                 for self.pkt_num, pkt in enumerate(conv_pkts):
-                    eth_frame = Ether(pkt[0])
+                    eth_frame = inet.Ether(pkt[0])
                     ip_pkt = eth_frame.payload
                     tcp_pkt = ip_pkt.payload
 
@@ -253,13 +256,13 @@ class EternalBlueExploit(BaseAttack.BaseAttack):
                         # TCP
                         tcp_pkt.setfieldval("sport", port_source)
                         # destination port is fixed 4444
-                        ## Window Size
+                        # Window Size
                         destination_origin_win = tcp_pkt.getfieldval("window")
                         if destination_origin_win not in destination_origin_wins:
                             destination_origin_wins[destination_origin_win] = destination_win_prob_dict.random()
                         new_win = destination_origin_wins[destination_origin_win]
                         tcp_pkt.setfieldval("window", new_win)
-                        ## MSS
+                        # MSS
                         tcp_options = tcp_pkt.getfieldval("options")
                         if tcp_options:
                             if tcp_options[0][0] == "MSS":
@@ -270,7 +273,8 @@ class EternalBlueExploit(BaseAttack.BaseAttack):
                         new_pkt.time = timestamp_next_pkt
 
                         pps = max(Util.get_interval_pps(complement_interval_pps, timestamp_next_pkt), 10)
-                        timestamp_next_pkt = Util.update_timestamp(timestamp_next_pkt, pps) + inter_arrival_times[self.pkt_num]# float(timeSteps.random())
+                        timestamp_next_pkt = Util.update_timestamp(timestamp_next_pkt, pps) + inter_arrival_times[
+                            self.pkt_num]  # float(timeSteps.random())
 
                     # Reply
                     else:
@@ -284,13 +288,13 @@ class EternalBlueExploit(BaseAttack.BaseAttack):
                         # TCP
                         tcp_pkt.setfieldval("dport", port_source)
                         # source port is fixed 4444
-                        ## Window Size
+                        # Window Size
                         source_origin_win = tcp_pkt.getfieldval("window")
                         if source_origin_win not in source_origin_wins:
                             source_origin_wins[source_origin_win] = source_win_prob_dict.random()
                         new_win = source_origin_wins[source_origin_win]
                         tcp_pkt.setfieldval("window", new_win)
-                        ## MSS
+                        # MSS
                         tcp_options = tcp_pkt.getfieldval("options")
                         if tcp_options:
                             if tcp_options[0][0] == "MSS":
@@ -300,7 +304,8 @@ class EternalBlueExploit(BaseAttack.BaseAttack):
                         new_pkt = (eth_frame / ip_pkt / tcp_pkt)
 
                         pps = max(Util.get_interval_pps(complement_interval_pps, timestamp_next_pkt), 10)
-                        timestamp_next_pkt = Util.update_timestamp(timestamp_next_pkt, pps) + inter_arrival_times[self.pkt_num]# float(timeSteps.random())
+                        timestamp_next_pkt = Util.update_timestamp(timestamp_next_pkt, pps) + inter_arrival_times[
+                            self.pkt_num]  # float(timeSteps.random())
 
                         new_pkt.time = timestamp_next_pkt