ID2T-toolkit-BotnetTraffic/code/Attack/MembersMgmtCommAttack.py

from random import randint, randrange, choice
from collections import deque

from Attack import BaseAttack
from Attack.AttackParameters import Parameter as Param
from Attack.AttackParameters import ParameterTypes

from ID2TLib import FileUtils
from ID2TLib.PacketGenerator import PacketGenerator
from ID2TLib.IPGenerator import MappingIPGenerator 
from ID2TLib.PcapAddressOperations import PcapAddressOperations
from ID2TLib.MapInputCSVToIDs import find_interval_with_most_comm
from ID2TLib.MacAddressGenerator import MacAddressGenerator
from ID2TLib.MessageType import MessageType
from ID2TLib.genRandPort import generateRandomPort
from ID2TLib import PaddingGenerator

class MembersMgmtCommAttack(BaseAttack.BaseAttack):
    def __init__(self):
        """
        Creates a new instance of the Membership Management Communication.

        """
        # Initialize communication
        super(MembersMgmtCommAttack, self).__init__("Membership Management Communication Attack (MembersMgmtCommAttack)", 
                                        "Injects Membership Management Communication", "Botnet communication")

        # Define allowed parameters and their type
        self.supported_params = {
            # parameters regarding attack 
            Param.INJECT_AT_TIMESTAMP: ParameterTypes.TYPE_FLOAT,
            Param.INJECT_AFTER_PACKET: ParameterTypes.TYPE_PACKET_POSITION,
            Param.PACKETS_PER_SECOND: ParameterTypes.TYPE_FLOAT,
            Param.PACKETS_LIMIT: ParameterTypes.TYPE_INTEGER_POSITIVE,
            Param.ATTACK_DURATION: ParameterTypes.TYPE_INTEGER_POSITIVE,

            # use num_attackers to specify number of communicating devices?
            Param.NUMBER_BOTS: ParameterTypes.TYPE_INTEGER_POSITIVE,

            # input file containing botnet communication
            Param.FILE_CSV: ParameterTypes.TYPE_FILEPATH,
            Param.FILE_XML: ParameterTypes.TYPE_FILEPATH,

            # the scope of communications
            Param.COMM_TYPE: ParameterTypes.TYPE_COMM_TYPE,

            # the percentage of IP reuse (if total and other is specified, percentages are multiplied)
            Param.IP_REUSE_TOTAL: ParameterTypes.TYPE_PERCENTAGE,
            Param.IP_REUSE_LOCAL: ParameterTypes.TYPE_PERCENTAGE,
            Param.IP_REUSE_EXTERNAL: ParameterTypes.TYPE_PERCENTAGE,

            # the user-selected padding to add to every packet
            Param.PACKET_PADDING: ParameterTypes.TYPE_PADDING
        }

    def init_params(self):
        """
        Initialize some parameters of this communication-attack using the user supplied command line parameters.
		The remaining parameters are implicitly set in the provided data file. Note: the timestamps in the file 
		have to be sorted in ascending order

        :param statistics: Reference to a statistics object.
        """
        # set class constants
        self.DEFAULT_XML_PATH = "resources/MembersMgmtComm_example.xml"
        # threshold for ID to be recognized as rather common in communication
        self.MOST_COMMON_THRES = 0.08

        # PARAMETERS: initialize with default values
        # (values are overwritten if user specifies them)
        self.add_param_value(Param.INJECT_AFTER_PACKET, randint(0, self.statistics.get_packet_count()))
        self.add_param_value(Param.PACKETS_PER_SECOND, 0)
        self.add_param_value(Param.FILE_XML, self.DEFAULT_XML_PATH)
        self.add_param_value(Param.ATTACK_DURATION, 100)
        self.add_param_value(Param.NUMBER_BOTS, 20)

        # default locality behavior
        self.add_param_value(Param.COMM_TYPE, "mixed")
        # TODO: change 1 to something better
        self.add_param_value(Param.IP_REUSE_TOTAL, 1)
        self.add_param_value(Param.IP_REUSE_LOCAL, 0.5)
        self.add_param_value(Param.IP_REUSE_EXTERNAL, 0.5)

        # add default additional padding
        self.add_param_value(Param.PACKET_PADDING, 0)

        
    def generate_attack_pcap(self):
        # parse input CSV or XML
        filepath_xml = self.get_param_value(Param.FILE_XML)
        filepath_csv = self.get_param_value(Param.FILE_CSV)

        # prefer XML input over CSV input
        if filepath_csv and filepath_xml == self.DEFAULT_XML_PATH:
            filepath_xml = FileUtils.parse_csv_to_xml(filepath_csv) 

        abstract_packets = FileUtils.parse_xml(filepath_xml)
        
        # find a good communication mapping
        duration = self.get_param_value(Param.ATTACK_DURATION)
        number_bots = self.get_param_value(Param.NUMBER_BOTS)
        comm_interval = find_interval_with_most_comm(abstract_packets, number_bots, duration)[0]

        if comm_interval is None:
            print("Error: There is no interval in the given CSV/XML that has enough communication")
            return 0, None

        mapped_ids, packet_start_idx, packet_end_idx = comm_interval["IDs"], comm_interval["Start"], comm_interval["End"]

        # determine most common communicating IDs
        total_mentions = 0
        most_common_thres = self.MOST_COMMON_THRES
        most_common_ids = {}
        for id_ in mapped_ids:
            total_mentions += mapped_ids[id_]

        for id_ in mapped_ids:
            count = mapped_ids[id_]
            mentions_percentage = count / total_mentions
            if mentions_percentage >= most_common_thres:
                most_common_ids[id_] = mentions_percentage

        # determine amount of reused IPs
        reuse_percent_total = self.get_param_value(Param.IP_REUSE_TOTAL)
        reuse_percent_external = self.get_param_value(Param.IP_REUSE_EXTERNAL)
        reuse_percent_local = self.get_param_value(Param.IP_REUSE_LOCAL)

        reuse_count_external = int(reuse_percent_total * reuse_percent_external * len(mapped_ids))
        reuse_count_local = int(reuse_percent_total * reuse_percent_local * len(mapped_ids))

        # create bot IP and MAC configs
        macgen = MacAddressGenerator() 
        comm_type = self.get_param_value(Param.COMM_TYPE)
        pcapops = PcapAddressOperations(self.statistics)
        bot_configs = {}

        # where IP is external, already set MAC to router MAC here
        if comm_type == "mixed":
            ...
        elif comm_type == "external":
            ...
        elif comm_type == "local":
            existing_local_ips = pcapops.get_existing_priv_ips(reuse_count_local)
            count_avail = min(reuse_count_local, len(existing_local_ips))
            ids = set(mapped_ids.keys())
            for ip in existing_local_ips:
                random_id = choice(tuple(ids))
                mac = self.statistics.process_db_query("macAddress(IPAddress=%s)" % ip)
                bot_configs[random_id] = {"Type": "local", "IP": ip, "MAC":mac}
                ids.remove(random_id)

            count_remaining = number_bots - count_avail
            new_local_ips = pcapops.get_new_priv_ips(count_remaining)
            for ip in new_local_ips:
                random_id = choice(tuple(ids))
                mac = macgen.random_mac()
                bot_configs[random_id] = {"Type": "local", "IP": ip, "MAC": mac}
                ids.remove(random_id)


        # create bot port configs
        for bot in bot_configs:
            bot_configs[bot]["Port"] = generateRandomPort()    

        if filepath_csv and filepath_xml == self.DEFAULT_XML_PATH:
            filepath_xml = FileUtils.parse_csv_to_xml(filepath_csv) 

        # Setup initial parameters for packet creation
        BUFFER_SIZE = 1000
        pkt_gen = PacketGenerator()
        file_timestamp_prv = float(abstract_packets[packet_start_idx]["Time"])
        pcap_timestamp = self.get_param_value(Param.INJECT_AT_TIMESTAMP)
        padding = self.get_param_value(Param.PACKET_PADDING)
        packets = deque(maxlen=BUFFER_SIZE)
        total_pkts = 0
        limit_packetcount = self.get_param_value(Param.PACKETS_LIMIT)
        limit_duration = duration
        duration = 0
        path_attack_pcap = None
        nl_requests = {}

        # create packets to write to pcap file
        for abst_packet in abstract_packets[packet_start_idx:packet_end_idx]:
            # map/retrieve addresses to ids from input file
            id_src, id_dst = abst_packet["Src"], abst_packet["Dst"]
            if (not id_src in bot_configs) or (not id_dst in bot_configs):
                continue

            ip_src, ip_dst = bot_configs[id_src]["IP"], bot_configs[id_dst]["IP"]
            mac_src, mac_dst = bot_configs[id_src]["MAC"], bot_configs[id_dst]["MAC"]
            port_src, port_dst = bot_configs[id_src]["Port"], bot_configs[id_dst]["Port"]

            # update timestamps and duration
            file_timestamp = float(abst_packet["Time"])
            file_time_delta = file_timestamp - file_timestamp_prv
            pcap_timestamp += file_time_delta
            duration += file_time_delta
            file_timestamp_prv = file_timestamp

            # if total number of packets has been sent or the attack duration has been exceeded, stop
            if ((limit_packetcount is not None and total_pkts >= limit_packetcount) or 
                    (limit_duration is not None and duration >= limit_duration)):
                break
        
            # create ip packet and add to packets list
            message_type = int(abst_packet["Type"])
            nl_size = 0

            if message_type == MessageType.SALITY_NL_REPLY.value:
                nl_size = randint(1, 25)

            packet = pkt_gen.generate_MMCOM_Packet(ip_src=ip_src, ip_dst=ip_dst, mac_src=mac_src, mac_dst=mac_dst, port_src=port_src, port_dst=port_dst, message_type=message_type, neighborlist_entries=nl_size)
            PaddingGenerator.add_padding(packet, padding)

            packet.time = pcap_timestamp
            packets.append(packet)
            total_pkts += 1

            # Store timestamp of first packet (for attack label)
            if total_pkts <= 1 :
                self.attack_start_utime = packets[0].time
            elif total_pkts % BUFFER_SIZE == 0: # every 1000 packets write them to the pcap file (append)
                packets = list(packets)
                PaddingGenerator.equal_length(packets)
                last_packet = packets[-1]
                path_attack_pcap = self.write_attack_pcap(packets, True, path_attack_pcap)
                packets = deque(maxlen=BUFFER_SIZE)

        if len(packets) > 0:
            packets = list(packets)
            PaddingGenerator.equal_length(packets)
            path_attack_pcap = self.write_attack_pcap(packets, True, path_attack_pcap)
            last_packet = packets[-1]

        # Store timestamp of last packet
        self.attack_end_utime = last_packet.time

        # Return packets sorted by packet by timestamp and total number of packets (sent)
        return total_pkts , path_attack_pcap