redesigned the ip-generator to use the new classes

code/ID2TLib/IPGenerator.py

@@ -1,28 +1,7 @@
 import random
-import re
+from . import IPv4 as ip
 
 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):
@@ -30,68 +9,40 @@ class IPGenerator:
 		self.generated_ips = set()
 		
 		if not include_private_ips:
-			for segment in self.PRIVATE_IP_SEGMENTS:
+			for segment in ip.ReservedIPBlocks.PRIVATE_IP_SEGMENTS:
 				self.add_to_blacklist(segment)
 		if not include_localhost:
-			self.add_to_blacklist(self.LOCALHOST_SEGMENT)
+			self.add_to_blacklist(ip.ReservedIPBlocks.LOCALHOST_SEGMENT)
 		if not include_multicast:
-			self.add_to_blacklist(self.MULTICAST_SEGMENT)
+			self.add_to_blacklist(ip.ReservedIPBlocks.MULTICAST_SEGMENT)
 		if not include_reserved:
-			self.add_to_blacklist(self.RESERVED_SEGMENT)
+			self.add_to_blacklist(ip.ReservedIPBlocks.RESERVED_SEGMENT)
 		if not include_link_local:
-			self.add_to_blacklist(self.ZERO_CONF_SEGMENT)
+			self.add_to_blacklist(ip.ReservedIPBlocks.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 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):
 		while True:
-			ip = random.randrange(0, 1 << 32)
+			random_ip = ip.IPAddress.from_int(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)
+			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()
 	
-	# 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
-		)
+	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):

code/ID2TLib/IPv4.py

@@ -108,6 +108,9 @@ class IPAddressBlock:
 		
 		return IPAddressBlock(ip, suffix)
 	
+	def block_size(self):
+		return 2 ** (32 - netmask)
+	
 	def _bitmask(self, netmask):
 		ones = lambda x: (1 << x) - 1