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

import random
from . import IPv4 as ip

class IPChooser:
	def random_ip(self):
		return ip.IPAddress.from_int(random.randrange(0, 1 << 32))
	
	def size(self):
		return 1 << 32
	
	def __len__(self):
		return self.size()

class IPChooserByRange(IPChooser):
	def __init__(self, ip_range):
		self.range = ip_range
	
	def random_ip(self):
		start = int(self.range.first_address())
		end = start + self.range.block_size()
		return ip.IPAddress.from_int(random.randrange(start, end))
	
	def size(self):
		return self.range.block_size()

class IPChooserByList(IPChooser):
	def __init__(self, ips):
		self.ips = list(ips)
		if not self.ips:
			raise ValueError("list of ips must not be empty")
	
	def random_ip(self):
		return random.choice(self.ips)
	
	def size(self):
		return len(self.ips)

class IPGenerator:
	def __init__(self, ip_chooser = IPChooser(), # include all ip-addresses by default (before the blacklist)
			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 ip.ReservedIPBlocks.PRIVATE_IP_SEGMENTS:
				self.add_to_blacklist(segment)
		if not include_localhost:
			self.add_to_blacklist(ip.ReservedIPBlocks.LOCALHOST_SEGMENT)
		if not include_multicast:
			self.add_to_blacklist(ip.ReservedIPBlocks.MULTICAST_SEGMENT)
		if not include_reserved:
			self.add_to_blacklist(ip.ReservedIPBlocks.RESERVED_SEGMENT)
		if not include_link_local:
			self.add_to_blacklist(ip.ReservedIPBlocks.ZERO_CONF_SEGMENT)
		if blacklist:
			for segment in blacklist:
				self.add_to_blacklist(segment)
		self.chooser = ip_chooser
	
	@staticmethod
	def from_range(range, *args, **kwargs):
		return IPGenerator(IPChooserByRange(range), *args, **kwargs)
	
	def add_to_blacklist(self, ip_segment):
		if isinstance(ip_segment, ip.IPAddressBlock):
			self.blacklist.append(ip_segment)
		else:
			self.blacklist.append(ip.IPAddressBlock.parse(ip_segment))
	
	def random_ip(self):
		if len(self.generated_ips) == self.chooser.size():
			raise ValueError("Exhausted the space of possible ip-addresses, no new unique ip-address can be generated")
		
		while True:
			random_ip = self.chooser.random_ip()
			
			if not self._is_in_blacklist(random_ip) and random_ip not in self.generated_ips:
				self.generated_ips.add(random_ip)
				return str(random_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()
	
	def _is_in_blacklist(self, ip: ip.IPAddress):
		return any(ip in block for block in self.blacklist)

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)