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

import re

class IPAddress:
	# 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)
	
	def __init__(self, intlist):
		if not isinstance(intlist, list) or not all(isinstance(n, int) for n in intlist):
			raise TypeError("The first constructor argument must be an list of ints")
		if not len(intlist) == 4 or not all(0 <= n <= 255 for n in intlist):
			raise ValueError("The integer list must contain 4 ints in range of 0 and 255, like an ip-address")
		
		self.ipnum = int.from_bytes(bytes(intlist), "big")
	
	@staticmethod
	def parse(ip: str):
		match = re.match("^" + IPAddress.IP_REGEXP + "$", ip)
		if not match:
			raise ValueError("%s is no ipv4-address" % ip)
		
		numbers = [int(match.group(i)) for i in range(1, 5)]
		return IPAddress(numbers)
	
	@staticmethod
	def from_int(numeric: int):
		if numeric not in range(1 << 32):
			raise ValueError("numeric value must be in uint-range")
		
		return IPAddress(list(numeric.to_bytes(4, "big")))
	
	def to_int(self):
		return self.ipnum
	
	def is_private(self):
		return ReservedIPBlocks.is_private(self)
	
	def is_localhost(self):
		return ReservedIPBlocks.is_localhost(self)
	
	def is_multicast(self):
		return ReservedIPBlocks.is_multicast(self)
	
	def is_reserved(self):
		return ReservedIPBlocks.is_reserved(self)
	
	def is_zero_conf(self):
		return ReservedIPBlocks.is_zero_conf(self)
	
	def _tuple(self):
		return tuple(self.ipnum.to_bytes(4, "big"))
	
	def __repr__(self):
		return "IPAddress([%i, %i, %i, %i])" % self._tuple()
	
	def __str__(self):
		return "%i.%i.%i.%i" % self._tuple()
	
	def __hash__(self):
		return self.ipnum
	
	def __eq__(self, other):
		if other is None:
			return False
		
		return isinstance(other, IPAddress) and self.ipnum == other.ipnum
	
	def __lt__(self, other):
		if other is None:
			raise TypeError("Cannot compare to None")
		if not isinstance(other, IPAddress):
			raise NotImplemented # maybe other can compare to self
		
		return self.ipnum < other.ipnum

class IPAddressBlock:
	CIDR_REGEXP = IPAddress.IP_REGEXP + r"(\/(3[0-2]|[12]?\d)|)?"
	
	def __init__(self, ip, netmask = 32):
		if isinstance(ip, str):
			ip = IPAddress.parse(ip)
		elif isinstance(ip, list):
			ip = IPAddress(ip)
		
		if not 1 <= netmask <= 32:
			raise ValueError("netmask must lie between 1 and 32")
		
		self.ipnum = ip.to_int() & self._bitmask(netmask)
		self.netmask = netmask
	
	@staticmethod
	def parse(cidr: str):
		match = re.match("^" + IPAddressBlock.CIDR_REGEXP + "$", cidr)
		if not match:
			raise ValueError("%s is no valid cidr-notation" % cidr)
		
		ip = [int(match.group(i)) for i in range(1, 5)]
		suffix = 32 if not match.group(6) else int(match.group(6))
		
		return IPAddressBlock(ip, suffix)
	
	def _bitmask(self, netmask):
		ones = lambda x: (1 << x) - 1
		
		return ones(32) ^ ones(32 - netmask)
	
	def __repr__(self):
		return "IPAddressBlock(%s, %i)" % (repr(IPAddress.from_int(self.ipnum)), self.netmask)
	
	def __self__(self):
		return IPAddress.from_int(self.ipnum) + "/" + str(self.netmask)
	
	def __contains__(self, ip):
		return (ip.to_int() & self._bitmask(self.netmask)) == self.ipnum

class ReservedIPBlocks:
	PRIVATE_IP_SEGMENTS = [
		IPAddressBlock.parse(block)
		for block in
		("10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16")
	]
	
	LOCALHOST_SEGMENT = IPAddressBlock.parse("127.0.0.0/8")
	
	MULTICAST_SEGMENT = IPAddressBlock.parse("224.0.0.0/4")
	RESERVED_SEGMENT = IPAddressBlock.parse("240.0.0.0/4")
	
	ZERO_CONF_SEGMENT = IPAddressBlock.parse("169.254.0.0/16")
	
	@staticmethod
	def is_private(ip):
		return any(ip in block for block in ReservedIPBlocks.PRIVATE_IP_SEGMENTS)
	
	@staticmethod
	def is_localhost(ip):
		return ip in ReservedIPBlocks.LOCALHOST_SEGMENT
	
	@staticmethod
	def is_multicast(ip):
		return ip in ReservedIPBlocks.MULTICAST_SEGMENT
	
	@staticmethod
	def is_reserved(ip):
		return ip in ReservedIPBlocks.RESERVED_SEGMENT
	
	@staticmethod
	def is_zero_conf(ip):
		return ip in ReservedIPBlocks.ZERO_CONF_SEGMENT