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

import random
import re

class IPGenerator:
	# see wikipedia
	PRIVATE_IP_SEGMENTS = [
		"10.0.0.0/8",
		"172.16.0.0/12",
		"192.168.0.0/16"
	]
	
	LOCALHOST_SEGMENT = "127.0.0.0/8"
	
	MULTICAST_SEGMENT = "224.0.0.0/4" # class D segment
	RESERVED_SEGMENT  = "240.0.0.0/4" # class E segment
	
	ZERO_CONF_SEGMENT = "169.254.0.0/16" # link local segment
	
	# a number between 0 and 255, no leading zeros
	_IP_NUMBER_REGEXP = r"(25[0-5]|2[0-4]\d|1?[1-9]?\d)"
	# 4 numbers between 0 and 255, joined together with dots
	IP_REGEXP = r"{0}\.{0}\.{0}\.{0}".format(_IP_NUMBER_REGEXP)
	# an ip address with an optional cidr-suffix
	CIDR_REGEXP = IP_REGEXP + r"(\/(3[0-2]|[12]?\d)|)?"
	
	def __init__(self, include_private_ips = False, include_localhost = False,
			include_multicast = False, include_reserved = False,
			include_link_local = False, blacklist = None):
		self.blacklist = []
		self.generated_ips = set()
		
		if not include_private_ips:
			for segment in self.PRIVATE_IP_SEGMENTS:
				self.add_to_blacklist(segment)
		if not include_localhost:
			self.add_to_blacklist(self.LOCALHOST_SEGMENT)
		if not include_multicast:
			self.add_to_blacklist(self.MULTICAST_SEGMENT)
		if not include_reserved:
			self.add_to_blacklist(self.RESERVED_SEGMENT)
		if not include_link_local:
			self.add_to_blacklist(self.ZERO_CONF_SEGMENT)
		if blacklist:
			for segment in blacklist:
				self.add_to_blacklist(segment)	
	
	def add_to_blacklist(self, ip_segment: str):
		self.blacklist.append(self._parse_cidr(ip_segment))
	
	def random_ip(self):
		while True:
			ip = random.randrange(0, 1 << 32)
			
			if not self._is_in_blacklist(ip) and ip not in self.generated_ips:
				self.generated_ips.add(ip)
				return self._ip_to_str(ip)
	
	def clear(self, clear_blacklist = True, clear_generated_ips = True):
		if clear_blacklist: self.blacklist.clear()
		if clear_generated_ips: self.generated_ips.clear()
	
	# parses a str in cidr-notation and returns a tuple with 2 elements
	# the first element is the ip-address as int, the second one is the cidr-suffix as int
	def _parse_cidr(self, ip_segment: str):
		match = re.match("^" + IPGenerator.CIDR_REGEXP + "$", ip_segment)
		if not match:
			raise ValueError("%s is no ip in cidr-notation" % ip_segment)
		
		ip = [int(match.group(i)) for i in range(1, 5)]
		suffix = 32 if not match.group(6) else int(match.group(6))
		
		numeric_ip = (ip[0] << 24) | (ip[1] << 16) | (ip[2] << 8) | ip[3]
		
		return (numeric_ip & self._netmask(suffix), suffix)
	
	def _is_in_blacklist(self, ip: int):
		for black_ip, cidr in self.blacklist:
			if (ip & self._netmask(cidr)) == black_ip:
				return True
		
		return False
	
	def _netmask(self, suffix: int):
		ones = lambda x: (1 << x) - 1
		
		return ones(32) ^ ones(32 - suffix)
	
	def _ip_to_str(self, ip: int):
		return "%i.%i.%i.%i" % (
			ip >> 24,
			(ip >> 16) & 255,
			(ip >> 8) & 255,
			ip & 255
		)

class MappingIPGenerator(IPGenerator):
	def __init__(self, *args, **kwargs):
		super().__init__(self, *args, **kwargs)
		
		self.mapping = {}
	
	def clear(self, clear_generated_ips = True, *args, **kwargs):
		super().clear(self, clear_generated_ips = clear_generated_ips, *args, **kwargs)
		if clear_generated_ips:
			self.mapping  = {}
	
	def get_mapped_ip(self, key):
		if key not in self.mapping:
			self.mapping[key] = self.random_ip()
		
		return self.mapping[key]
	
	def __getitem__(self, item):
		return self.get_mapped_ip(item)