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

from enum import Enum
class MessageType(Enum):
    """
    Defines possible Message types
    """

    TIMEOUT = 3
    SALITY_NL_REQUEST = 101
    SALITY_NL_REPLY = 102
    SALITY_HELLO = 103
    SALITY_HELLO_REPLY = 104

from random import randint, randrange, choice
from collections import deque
from scipy.stats import gamma
from lea import Lea

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

from ID2TLib import FileUtils, PaddingGenerator
from ID2TLib.PacketGenerator import PacketGenerator
from ID2TLib.IPGenerator import IPGenerator 
from ID2TLib.PcapAddressOperations import PcapAddressOperations
from ID2TLib.MapInputCSVToIDs import find_interval_with_most_comm, determine_id_roles
from ID2TLib.MacAddressGenerator import MacAddressGenerator
from ID2TLib.PortGenerator import gen_random_server_port

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
        }

        # create dict with MessageType values for fast name lookup
        self.msg_types = {}
        for msg_type in MessageType:
            self.msg_types[msg_type.value] = msg_type

    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
        # probability for initiator ID to be local
        self.PROB_INIT_IS_LOCAL = 0.8
        # probability for responder ID to be local if comm_type is mixed
        self.PROB_RESPND_IS_LOCAL = 0.2

        # PARAMETERS: initialize with default values
        # (values are overwritten if user specifies them)
        # print(self.statistics.get_packet_count())
        self.add_param_value(Param.INJECT_AFTER_PACKET, randint(0, self.statistics.get_packet_count()))
        # print(self.get_param_value(Param.INJECT_AT_TIMESTAMP))
        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):
        def add_ids_to_config(ids_to_add, existing_ips, new_ips, bot_configs, idtype="local", router_mac=""):
            ids = ids_to_add.copy()
            macgen = MacAddressGenerator()
            for ip in existing_ips:
                random_id = choice(tuple(ids))
                mac = self.statistics.process_db_query("macAddress(IPAddress=%s)" % ip)
                bot_configs[random_id] = {"Type": idtype, "IP": ip, "MAC": mac}
                ids.remove(random_id)

            for ip in new_ips:
                random_id = choice(tuple(ids))
                if idtype == "local":
                    mac = macgen.random_mac()
                elif idtype == "external":
                    mac = router_mac
                bot_configs[random_id] = {"Type": idtype, "IP": ip, "MAC": mac}
                ids.remove(random_id)

        def index_increment(number: int, max: int):
            if number + 1 < max:
                return number + 1
            else:
                return 0

        # 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, id_comms, packet_start_idx, packet_end_idx = comm_interval["IDs"], comm_interval["Comms"], comm_interval["Start"], comm_interval["End"]

        # print start and end time of mapped interval
        # print(abstract_packets[packet_start_idx]["Time"])
        # print(abstract_packets[packet_end_idx]["Time"])

        # 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
        ipgen = IPGenerator()
        comm_type = self.get_param_value(Param.COMM_TYPE)
        pcapops = PcapAddressOperations(self.statistics)
        router_mac = pcapops.get_probable_router_mac()
        bot_configs = {}
        external_ids = set()
        local_ids = set()
        # print(self.get_param_value(Param.INJECT_AT_TIMESTAMP))
        if comm_type == "local":
            local_ids = set(mapped_ids.keys())
        else:
            init_local_or_external = Lea.fromValFreqsDict({"local": self.PROB_INIT_IS_LOCAL*100, "external": (1-self.PROB_INIT_IS_LOCAL)*100})
            mixed_respnd_is_local = Lea.fromValFreqsDict({"local": self.PROB_RESPND_IS_LOCAL*100, "external": (1-self.PROB_RESPND_IS_LOCAL)*100})
            ids = set(mapped_ids.keys())
            init_ids, respnd_ids, both_ids = determine_id_roles(abstract_packets[packet_start_idx:packet_end_idx+1], ids)

            # assign IDs in 'both' local everytime for mixed? 
            initiators = list(init_ids) + list(both_ids)
            for id_ in initiators:
                if id_ in local_ids or id_ in external_ids:
                    continue

                pos = init_local_or_external.random()          
                if pos == "local":
                    local_ids.add(id_) 
                    for id_comm in id_comms:
                        ids = id_comm.split("-")
                        other = ids[0] if id_ == ids[1] else ids[1] 
                        
                        # what if before other was external ...
                        if other in local_ids or other in external_ids:
                            continue 

                        if comm_type == "mixed":
                            other_pos = mixed_respnd_is_local.random()
                            if other_pos == "local":
                                local_ids.add(other)
                            elif other_pos == "external":
                                external_ids.add(other)
                        elif comm_type == "external":
                            if not other in initiators:
                                external_ids.add(other)

                elif pos == "external":
                    external_ids.add(id_)
                    for id_comm in id_comms:
                        ids = id_comm.split("-")
                        other = ids[0] if id_ == ids[1] else ids[1] 
                        
                        # what if before other was external ...
                        if other in local_ids or other in external_ids:
                            continue 
                        if not other in initiators:
                            local_ids.add(other)


        # print(initiators)
        # print(local_ids)
        # print(external_ids)

        number_local_ids, number_external_ids = len(local_ids), len(external_ids)
        if number_local_ids > 0:
            reuse_count_local = int(reuse_percent_total * reuse_percent_local * number_local_ids) 
            existing_local_ips = pcapops.get_existing_priv_ips(reuse_count_local)
            new_local_ips = pcapops.get_new_priv_ips(number_local_ids - len(existing_local_ips))
            add_ids_to_config(local_ids, existing_local_ips, new_local_ips, bot_configs)

        if number_external_ids > 0:
            reuse_count_external = int(reuse_percent_total * reuse_percent_external * number_external_ids) 
            existing_external_ips = pcapops.get_existing_external_ips(reuse_count_external)
            remaining = len(external_ids) - len(existing_external_ips)
            new_external_ips = [ipgen.random_ip() for _ in range(remaining)]
            add_ids_to_config(external_ids, existing_external_ips, new_external_ips, bot_configs, idtype="external", router_mac=router_mac)

        # print(bot_configs)

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

        # create realistic ttl for every bot
        # Gamma distribution parameters derived from MAWI 13.8G dataset
        ids = bot_configs.keys()
        alpha, loc, beta = (2.3261710235, -0.188306914406, 44.4853123884)
        gd = gamma.rvs(alpha, loc=loc, scale=beta, size=len(ids))

        for pos, bot in enumerate(bot_configs):
            is_invalid = True
            pos_max = len(gd)
            while is_invalid:
                ttl = int(round(gd[pos]))
                if 0 < ttl < 256:  # validity check
                    is_invalid = False
                else:
                    pos = index_increment(pos, pos_max)
            bot_configs[bot]["TTL"] = ttl

        # 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+1]:
            # 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"]
            ttl = bot_configs[id_src]["TTL"]
            type_src, type_dst = bot_configs[id_src]["Type"], bot_configs[id_dst]["Type"]

            if type_src == "external" and type_dst == "external":
                continue
            if comm_type == "external":
                if type_src == "local" and type_dst == "local":
                    continue


            # 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 = self.msg_types[int(abst_packet["Type"])]
            nl_size = 0

            if message_type == MessageType.SALITY_NL_REPLY:
                nl_size = randint(1, 25)
            elif message_type == MessageType.TIMEOUT:
                continue

            packet = pkt_gen.generate_mmcom_packet(ip_src=ip_src, ip_dst=ip_dst, ttl=ttl, 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