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

from random import choice

from ID2TLib import Statistics
from ID2TLib.IPv4 import IPAddress

is_ipv4 = IPAddress.is_ipv4

class PcapAddressOperations():

    def __init__(self, statistics: Statistics, uncertain_ip_mult: int=3):
        """
        Initializes a pcap information extractor that uses the provided statistics for its operations.

        :param statistics: The statistics of the pcap file
        :param uncertain_ip_mult: the mutliplier to create new address space when the remaining observed space has been drained
        """
        self.statistics = statistics
        self.UNCERTAIN_IPSPACE_MULTIPLIER = uncertain_ip_mult
        self._init_ipaddress_ops()

    def get_probable_router_mac(self):
        """
        Returns the most probable router MAC address based on the most used MAC address in the statistics.
        """
        self.probable_router_mac, count = self.statistics.process_db_query("most_used(macAddress)", print_results=False)[0]
        return self.probable_router_mac     # and count as a measure of certainty?

    def pcap_contains_priv_ips(self):
        """
        Returns True if the provided traffic contains private IPs, otherwise False.
        """
        return self.contains_priv_ips

    def get_priv_address_range(self):
        """
        Returns a tuple with the start and end of the observed private IP range.
        """
        # better way to handle error?
        if not self.pcap_contains_priv_ips():
            print("Error: .pcap does not contain any private ips.")
            return -1, -1
        return str(self.min_priv_ip), str(self.max_priv_ip)

    def get_count_rem_priv_ips(self):
        """
        Returns the number of private IPs in the pcap file that have not aldready been returned by get_existing_priv_ips.
        """
        return len(self.remaining_priv_ips)

    def get_existing_priv_ips(self, count: int=1):
        """
        Returns the given number of private IPs that are existent in the pcap file.

        :param count: the number of IPs to return
        :return: the chosen private IPs
        """

        # reasonable to include this?
        if not self.pcap_contains_priv_ips():
            print("Warning: .pcap does not contain any private ips.")
            return []

        if count > len(self.priv_ips):
            print("Warning: There are no {} priv IPs in the .pcap file. Returning all existing priv IPs.".format(count))

        total = min(len(self.priv_ips), count)

        retr_priv_ips = []
        priv_ips = self.remaining_priv_ips
        for _ in range(0, total):
            random_priv_ip = choice(tuple(priv_ips))
            retr_priv_ips.append(str(random_priv_ip))
            priv_ips.remove(random_priv_ip)

        return retr_priv_ips

    # also use IPs below minimum observed IP?
    # offset for later, start at x after minimum? e.g. start at 192.168.0.100
    # exclude the last IP of an IP segment because its broadcast?
    def get_new_priv_ips(self, count: int=1):
        """
        Returns in the pcap not existent private IPs that match the used segment. IPs can be returned
        that are either between the minimum and maximum observed IP and are therefore considered certain
        or that are above the observed maximum address, are more likely to not belong to the network using the
        private IP segment and are therefore considered uncertain.

        :param count: the number of IPs to return
        :return: the newly created private IP addresses
        """

        if not self.pcap_contains_priv_ips():
            print("Error: .pcap does not contain any private ips.")
            return []

        unused_priv_ips = self.unused_priv_ips
        uncertain_priv_ips = self.uncertain_priv_ips

        # warning reasonable?
        if count > len(unused_priv_ips):
            print("Warning: there are no {0} unused certain priv IPs in the .pcap file.\n \
                Returning {1} certain and {2} uncertain priv IPs.".format(count, len(unused_priv_ips), count-len(unused_priv_ips)))

        count_certain = min(count, len(unused_priv_ips))
    
        retr_priv_ips = []
        for _ in range(0, count_certain):
            random_priv_ip = choice(tuple(unused_priv_ips))
            retr_priv_ips.append(str(random_priv_ip))
            unused_priv_ips.remove(random_priv_ip)

        # retrieve uncertain priv ips
        if count_certain < count:
            count_uncertain = count - count_certain

            # check if new uncertain IPs have to be created
            if len(uncertain_priv_ips) < count_uncertain:
                ipspace_multiplier = self.UNCERTAIN_IPSPACE_MULTIPLIER

                max_new_ip = self.max_uncertain_priv_ip.to_int() + ipspace_multiplier * count_uncertain
                # adjust IP space multiplier and prevent broadcast address of private segment from being chosen as new IP
                while max_new_ip >= self.priv_ip_segment.last_address().to_int():
                    max_new_ip -= 1
            
                count_new_ips = max_new_ip - self.max_uncertain_priv_ip.to_int()
                if count_new_ips < count_uncertain:
                    print("Error: Cannot generate enough new private IPs because they would exceed the maximum private segment IP. Returning {}.".format(count_new_ips))

                # create ipspace_multiplier * count_uncertain new uncertain IP addresses
                last_gen_ip = None
                for i in range(1, count_new_ips + 1):
                    ip = IPAddress.from_int(self.max_uncertain_priv_ip.to_int() + i)
                    # exclude the broadcast address
                    if ip.to_int() >= self.priv_ip_segment.last_address().to_int():
                        break
                    uncertain_priv_ips.add(ip)
                    last_gen_ip = ip
                self.max_uncertain_priv_ip = last_gen_ip

            # choose the uncertain IPs to return
            total_uncertain = min(count_uncertain, len(uncertain_priv_ips))
            for _ in range(0, total_uncertain):
                random_priv_ip = choice(tuple(uncertain_priv_ips))
                retr_priv_ips.append(str(random_priv_ip))
                uncertain_priv_ips.remove(random_priv_ip)

        return retr_priv_ips

    def _init_ipaddress_ops(self):
        """
        Load and process data needed to perform functions on the IP addresses contained in the statistics
        """

        all_ips = self.statistics.process_db_query("all(ipAddress)", print_results=False)

        # find the private IP segment in use
        priv_ip_segment = None
        self.contains_priv_ips = False
        first_priv_ip = None
        first_priv_ip_idx = -1

        # for that iterate over all IPs until the first private IP is found
        for i, ip in enumerate(all_ips):
            if not is_ipv4(ip): 
                continue
            ip = IPAddress.parse(ip)

            if ip.is_private():
                priv_ip_segment = ip.get_private_segment()
                first_priv_ip_idx = i
                first_priv_ip = ip
                self.contains_priv_ips = True
                break

        if not self.contains_priv_ips:
            #print("The Pcap File does not contain any private IPs")
            return

        # get minimum and maximum seen private IP. The addresses in-bewteen are considered 
        # as certain to be part of the network the pcap traffic is from
        min_priv_ip, max_priv_ip = first_priv_ip, first_priv_ip
        priv_ips = {first_priv_ip}
        for ip in all_ips[first_priv_ip_idx+1:]:
            if not is_ipv4(ip): 
                continue
            ip = IPAddress.parse(ip)

            if ip in priv_ip_segment:
                priv_ips.add(ip)
                if ip > max_priv_ip:
                    max_priv_ip = ip
                elif ip < min_priv_ip:
                    min_priv_ip = ip

        # save the certain unused priv IPs of the network
        unused_priv_ips = set()
        for i in range (min_priv_ip.to_int() + 1, max_priv_ip.to_int()):
            ip = IPAddress.from_int(i)
            if not ip in priv_ips:
                unused_priv_ips.add(ip)

        # save the gathered information for efficient later use
        self.min_priv_ip, self.max_priv_ip = min_priv_ip, max_priv_ip
        self.max_uncertain_priv_ip = max_priv_ip
        self.priv_ips = frozenset(priv_ips)
        self.remaining_priv_ips = priv_ips
        self.unused_priv_ips = unused_priv_ips
        self.generated_uncertain_ips = set()
        self.uncertain_priv_ips = set()
        self.priv_ip_segment = priv_ip_segment