Added the possibility to create ip-addresses from user-defined blocks and lists

code/ID2TLib/IPGenerator.py

@@ -1,8 +1,43 @@
 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, include_private_ips = False, include_localhost = False,
+	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 = []
@@ -21,7 +56,12 @@ class IPGenerator:
 			self.add_to_blacklist(ip.ReservedIPBlocks.ZERO_CONF_SEGMENT)
 		if blacklist:
 			for segment in blacklist:
-				self.add_to_blacklist(segment)	
+				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):
@@ -30,8 +70,11 @@ class IPGenerator:
 			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 = ip.IPAddress.from_int(random.randrange(0, 1 << 32))
+			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)

code/ID2TLib/IPv4.py

@@ -81,6 +81,9 @@ class IPAddress:
 			raise NotImplemented # maybe other can compare to self
 		
 		return self.ipnum < other.ipnum
+	
+	def __int__(self):
+		return self.ipnum
 
 class IPAddressBlock:
 	CIDR_REGEXP = IPAddress.IP_REGEXP + r"(\/(3[0-2]|[12]?\d)|)?"