Bläddra i källkod

Added an ip-generator, tests and further documentation will be postponed

Denis Waßmann 7 år sedan
förälder
incheckning
7f01aa5cee
1 ändrade filer med 114 tillägg och 0 borttagningar
  1. 114 0
      code/ID2TLib/IPGenerator.py

+ 114 - 0
code/ID2TLib/IPGenerator.py

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