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

from random import randint
from collections import deque

from Attack import BaseAttack
from Attack.AttackParameters import Parameter as Param
from Attack.AttackParameters import ParameterTypes
from ID2TLib import IPPacketGenerator as IPPktGen, FileUtils
from ID2TLib.IPGenerator import MappingIPGenerator 


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_ATTACKERS: ParameterTypes.TYPE_INTEGER_POSITIVE,

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

    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.MESSAGE_TYPES = {2: "SALITY", 3: "TIMEOUT"}
        self.DEFAULT_XML_PATH = "resources/MembersMgmtComm_example.xml"

        # 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)
        
    def generate_attack_pcap(self):
        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) 

        comm_entries = FileUtils.parse_xml(filepath_xml)

        # Setup initial parameters for packet creation
        BUFFER_SIZE = 1000
        id_to_ip_mapper = MappingIPGenerator()
        file_timestamp_prv = float(comm_entries[0]["Time"])
        pcap_timestamp = self.get_param_value(Param.INJECT_AT_TIMESTAMP)
        duration = 0
        packets = deque(maxlen=BUFFER_SIZE)
        total_pkts = 0
        limit_packetcount = self.get_param_value(Param.PACKETS_LIMIT)
        limit_duration = self.get_param_value(Param.ATTACK_DURATION)
        path_attack_pcap = None

        # create packets to write to pcap file
        for entry in comm_entries:
            # map/retrieve ip addresses to ids from input file
            ip_src = id_to_ip_mapper.get_mapped_ip(entry["Src"])
            ip_dst = id_to_ip_mapper.get_mapped_ip(entry["Dst"])
            
            # update timestamps and duration
            file_timestamp = float(entry["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 = self.MESSAGE_TYPES[int(entry["Type"])]
            packet = IPPktGen.generate_ip_packet(ip_src=ip_src, ip_dst=ip_dst, payload=message_type)
            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)
                last_packet = packets[-1]
                #packets = sorted(packets, key=lambda pkt: pkt.time)	# not necessary yet
                path_attack_pcap = self.write_attack_pcap(packets, True, path_attack_pcap)
                packets = []

        if len(packets) > 0:
            #packets = sorted(packets, key=lambda pkt: pkt.time)	# not necessary yet
            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