ID2T-toolkit-BotnetTraffic/code/ID2TLib/MapInputCSVToIDs.py

from Attack.MembersMgmtCommAttack import MessageType

# needed because of machine inprecision. E.g A time difference of 0.1s is stored as >0.1s
EPS_TOLERANCE = 1e-13  # works for a difference of 0.1, no less

def find_interval_with_most_comm(packets: list, number_ids: int, max_int_time: float):
    """
    Finds a time interval of the given seconds where the given number of ids communicate among themselves the most.
    
    :param packets: The packets containing the communication
    :param number_ids: The number of ids that are to be considered
    :param max_int_time: A short description of the attack.
    :return: A triple consisting of the ids, as well as start and end idx with respect to the given packets. 
    """

    def get_nez_msg_counts(msg_counts: dict):
        """
        Filters out all msg_counts that have 0 as value
        """
        nez_msg_counts = dict()
        for msg in msg_counts.keys():
            count = msg_counts[msg]
            if count > 0:
                nez_msg_counts[msg] = count
        return nez_msg_counts

    def greater_than(a: float, b: float):
        """
        A greater than operator desgined to handle slight machine inprecision up to EPS_TOLERANCE.
        :return: True if a > b, otherwise False
        """
        return b - a < -EPS_TOLERANCE

    def change_msg_counts(msg_counts: dict, idx: int, add=True):
        """
        Changes the value of the message count of the message occuring in the packet specified by the given index. 
        Adds 1 if add is True and subtracts 1 otherwise.
        """
        change = 1 if add else -1
        id_src, id_dst = packets[idx]["Src"], packets[idx]["Dst"]
        src_to_dst = "{0}-{1}".format(id_src, id_dst)
        dst_to_src = "{0}-{1}".format(id_dst, id_src)

        if src_to_dst in msg_counts.keys():
            msg_counts[src_to_dst] += change
        elif dst_to_src in msg_counts.keys():
            msg_counts[dst_to_src] += change
        elif add:
            msg_counts[src_to_dst] = 1

    def count_ids_in_msg_counts(msg_counts: dict):
        """
        Counts all ids that are involved in messages with a non zero message count
        """
        ids = set()
        for msg in msg_counts.keys():
            src, dst = msg.split("-")
            ids.add(dst)
            ids.add(src)
        return len(ids)

    def get_msg_count_first_ids(msg_counts: list):
        """
        Finds the ids that communicate among themselves the most with respect to the given message counts.
        :param msg_counts: a sorted list of message counts where each entry is a tuple of key and value
        :return: The picked ids and their total message count as a tuple
        """
        # if order of most messages is important, use an additional list
        picked_ids = set()
        total_msg_count = 0

        # iterate over every message count
        for i, msg in enumerate(msg_counts):
            count_picked_ids = len(picked_ids)
            id_one, id_two = msg[0].split("-")

            # if enough ids have been found, stop
            if count_picked_ids >= number_ids:
                break

            # if two ids can be added without exceeding the desired number of ids, add them
            if count_picked_ids - 2 <= number_ids:
                picked_ids.add(id_one)
                picked_ids.add(id_two)
                total_msg_count += msg[1]

            # if there is only room for one more id to be added, 
            # find one that is already contained in the picked ids
            else:
                for j, msg in enumerate(msg_counts[i:]):
                    id_one, id_two = msg[0].split("-")
                    if id_one in picked_ids:
                        picked_ids.add(id_two)
                        total_msg_count += msg[1]
                        break
                    elif id_two in picked_ids:
                        picked_ids.add(id_one)
                        total_msg_count += msg[1]
                        break
                break

        return picked_ids, total_msg_count

    def get_indv_id_counts_and_comms(picked_ids: dict, msg_counts: dict):
        """ 
        Retrieves the total mentions of one ID in the communication pattern
        and all communication entries that include only picked IDs.
        """
        indv_id_counts = {}
        id_comms = set()
        for msg in msg_counts:
            ids = msg.split("-")
            if ids[0] in picked_ids and ids[1] in picked_ids:
                msg_other_dir = "{}-{}".format(ids[1], ids[0])
                if (not msg in id_comms) and (not msg_other_dir in id_comms):
                    id_comms.add(msg)

                for id_ in ids:
                    if id_ in indv_id_counts:
                        indv_id_counts[id_] += msg_counts[msg]
                    else:
                        indv_id_counts[id_] = msg_counts[msg]

        return indv_id_counts, id_comms


    # first find all possible intervals that contain enough IDs that communicate among themselves
    idx_low, idx_high = 0, 0
    msg_counts = dict()
    possible_intervals = []

    # Iterate over all packets from start to finish and process the info of each packet
    # If time of packet within time interval, update the message count for this communication
    # If time of packet exceeds time interval, substract from the message count for this communication
    while True:
        if idx_high < len(packets):
            cur_int_time = float(packets[idx_high]["Time"]) - float(packets[idx_low]["Time"])
 
        # if current interval time exceeds time interval, save the message counts if appropriate, or stop if no more packets
        if greater_than(cur_int_time, max_int_time) or idx_high >= len(packets):
            # get all message counts for communications that took place in the current intervall
            nez_msg_counts = get_nez_msg_counts(msg_counts)

            # if we have enough ids as specified by the caller, mark as possible interval
            if count_ids_in_msg_counts(nez_msg_counts) >= number_ids:
                # possible_intervals.append((nez_msg_counts, packets[idx_low]["Time"], packets[idx_high-1]["Time"]))
                possible_intervals.append((nez_msg_counts, idx_low, idx_high - 1))

            if idx_high >= len(packets):
                break

        # let idx_low "catch up" so that the current interval time fits into the interval time specified by the caller
        while greater_than(cur_int_time, max_int_time):
            change_msg_counts(msg_counts, idx_low, add=False)
            idx_low += 1
            cur_int_time = float(packets[idx_high]["Time"]) - float(packets[idx_low]["Time"])

        # consume the new packet at idx_high and process its information
        change_msg_counts(msg_counts, idx_high)
        idx_high += 1


    # now find the interval in which as many ids as specified communicate the most in the given time interval
    summed_intervals = []
    sum_intervals_idxs = []
    cur_highest_sum = 0

    # for every interval compute the sum of msg_counts of the first most communicative ids and eventually find
    # the interval(s) with most communication and its ids
    for j, interval in enumerate(possible_intervals):
        msg_counts = interval[0].items()
        sorted_msg_counts = sorted(msg_counts, key=lambda x: x[1], reverse=True)
        picked_ids, msg_sum = get_msg_count_first_ids(sorted_msg_counts)

        if msg_sum == cur_highest_sum:
            summed_intervals.append({"IDs": picked_ids, "MsgSum": msg_sum, "Start": interval[1], "End": interval[2]})
            sum_intervals_idxs.append(j)
        elif msg_sum > cur_highest_sum:
            summed_intervals = []
            sum_intervals_idxs = [j]
            summed_intervals.append({"IDs": picked_ids, "MsgSum": msg_sum, "Start": interval[1], "End": interval[2]})
            cur_highest_sum = msg_sum

    for j, interval in enumerate(summed_intervals):
        idx = sum_intervals_idxs[j]
        msg_counts_picked = possible_intervals[idx][0]
        indv_id_counts, id_comms = get_indv_id_counts_and_comms(interval["IDs"], msg_counts_picked)
        interval["IDs"] = indv_id_counts
        interval["Comms"] = id_comms

    return summed_intervals


def determine_id_roles(packets: list, all_ids: set):
    """
    Determine the role of every mapped ID. The role can be initiator, responder or both.
    :param packets: the mapped section of abstract packets
    :param all_ids: all IDs that were mapped/chosen
    :return: a dict that for every ID contains its role
    """
    init_ids, respnd_ids, both_ids = set(), set(), set()
    def process_initiator(id_: str):
        if id_ in both_ids:
            pass
        elif not id_ in respnd_ids:
            init_ids.add(id_)
        elif id_ in respnd_ids:
            respnd_ids.remove(id_)
            both_ids.add(id_)

    def process_responder(id_: str):
        if id_ in both_ids:
            pass
        elif not id_ in init_ids:
            respnd_ids.add(id_)
        elif id_ in init_ids:
            init_ids.remove(id_)
            both_ids.add(id_)

    for packet in packets:
        id_src, id_dst, msg_type = packet["Src"], packet["Dst"], packet["Type"]
        if (not id_src in all_ids) or (not id_dst in all_ids):
            continue
        if int(msg_type) in {MessageType.SALITY_HELLO.value, MessageType.SALITY_NL_REQUEST.value}:
            process_initiator(id_src)
            process_responder(id_dst)
        elif int(msg_type) in {MessageType.SALITY_HELLO_REPLY.value, MessageType.SALITY_NL_REPLY.value}:
            process_initiator(id_dst)
            process_responder(id_src)

    return init_ids, respnd_ids, both_ids