Add functionality to get private IPs from Pcap file

code/ID2TLib/PcapAddressOperations.py

@@ -0,0 +1,211 @@
+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
+
+