Merge branch 'develop' of https://git.tk.informatik.tu-darmstadt.de/leon.boeck/ID2T-toolkit-BotnetTraffic into develop

code/ID2TLib/IPv4.py

@@ -1,25 +1,42 @@
 import re
 
+
 class IPAddress:
+	"""
+	A simple class encapsulating an ip-address. An IPAddress can be constructed by string, int and 4-element-list
+	(e.g. [8, 8, 8, 8]). This is a leightweight class as it only contains string-to-ip-and-reverse-conversion
+	and some convenience methods.
+	"""
+	
 	# a number between 0 and 255, no leading zeros
 	_IP_NUMBER_REGEXP = r"(25[0-5]|2[0-4]\d|1\d\d|[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):
+		"""
+		Construct an ipv4-address with a list of 4 integers, e.g. to construct the ip 10.0.0.0 pass [10, 0, 0, 0]
+		"""
 		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")
 		
+		# For easier calculations store the ip as integer, e.g. 10.0.0.0 is 0x0a000000
 		self.ipnum = int.from_bytes(bytes(intlist), "big")
 	
 	@staticmethod
 	def parse(ip: str):
+		"""
+		Parse an ip-address-string. If the string does not comply to the ipv4-format a ValueError is raised
+		:param ip: A string-representation of an ip-address, e.g. "10.0.0.0"
+		:return: IPAddress-object describing the ip-address
+		"""
 		match = re.match("^" + IPAddress.IP_REGEXP + "$", ip)
 		if not match:
 			raise ValueError("%s is no ipv4-address" % ip)
 		
+		# the matches we get are the numbers of the ip-address (match 0 is the whole ip-address)
 		numbers = [int(match.group(i)) for i in range(1, 5)]
 		return IPAddress(numbers)
 	
@@ -28,41 +45,75 @@ class IPAddress:
 		if numeric not in range(1 << 32):
 			raise ValueError("numeric value must be in uint-range")
 		
+		# to_bytes is the easiest way to split a 32-bit int into bytes
 		return IPAddress(list(numeric.to_bytes(4, "big")))
 	
 	@staticmethod
 	def is_ipv4(ip: str):
+		"""
+		Check if the supplied string is in ipv4-format
+		"""
+		
 		match = re.match("^" + IPAddress.IP_REGEXP + "$", ip)
 		return True if match else False
 
 	def to_int(self):
+		"""
+		Convert the ip-address to a 32-bit uint, e.g. IPAddress.parse("10.0.0.255").to_int() returns 0x0a0000ff
+		"""
 		return self.ipnum
 	
 	def is_private(self):
+		"""
+		Returns a boolean indicating if the ip-address lies in the private ip-segments (see ReservedIPBlocks)
+		"""
 		return ReservedIPBlocks.is_private(self)
 	
 	def get_private_segment(self):
+		"""
+		Return the private ip-segment the ip-address belongs to (there are several)
+		If this ip does not belong to a private ip-segment a ValueError is raised
+		:return: IPAddressBlock
+		"""
 		return ReservedIPBlocks.get_private_segment(self)
 
 	def is_localhost(self):
+		"""
+		Returns a boolean indicating if the ip-address lies in the localhost-segment
+		"""
 		return ReservedIPBlocks.is_localhost(self)
 	
 	def is_multicast(self):
+		"""
+		Returns a boolean indicating if the ip-address lies in the multicast-segment
+		"""
 		return ReservedIPBlocks.is_multicast(self)
 	
 	def is_reserved(self):
+		"""
+		Returns a boolean indicating if the ip-address lies in the reserved-segment
+		"""
 		return ReservedIPBlocks.is_reserved(self)
 	
 	def is_zero_conf(self):
+		"""
+		Returns a boolean indicating if the ip-address lies in the zeroconf-segment
+		"""
 		return ReservedIPBlocks.is_zero_conf(self)
 	
 	def _tuple(self):
 		return tuple(self.ipnum.to_bytes(4, "big"))
 	
 	def __repr__(self):
+		"""
+		Following the python style guide, eval(repr(obj)) should equal obj
+		"""
 		return "IPAddress([%i, %i, %i, %i])" % self._tuple()
 	
 	def __str__(self):
+		"""
+		Return the ip-address described by this object in ipv4-format
+		"""
 		return "%i.%i.%i.%i" % self._tuple()
 	
 	def __hash__(self):
@@ -86,9 +137,21 @@ class IPAddress:
 		return self.ipnum
 
 class IPAddressBlock:
+	"""
+	This class describes a block of IPv4-addresses, just as a string in CIDR-notation does.
+	It can be seen as a range of ip-addresses. To check if a block contains a ip-address
+	simply use "ip in ip_block"
+	"""
+	
+	# this regex describes CIDR-notation (an ip-address plus "/XX", whereas XX is a number between 1 and 32)
 	CIDR_REGEXP = IPAddress.IP_REGEXP + r"(\/(3[0-2]|[12]?\d)|)?"
 	
 	def __init__(self, ip, netmask = 32):
+		"""
+		Construct a ip-block given a ip-address and a netmask. Given an ip and a netmask,
+		the constructed ip-block will describe the range ip/netmask (e.g. 127.0.0.1/8)
+		:param ip: An ip-address, represented as IPAddress, string or 4-element-list
+		"""
 		if isinstance(ip, str):
 			ip = IPAddress.parse(ip)
 		elif isinstance(ip, list):
@@ -97,12 +160,17 @@ class IPAddressBlock:
 		if not 1 <= netmask <= 32:
 			raise ValueError("netmask must lie between 1 and 32")
 		
+		# clear the unnecessary bits in the base-ip, e.g. this will convert 10.0.0.255/24 to 10.0.0.0/24 which are equivalent
 		self.ipnum = ip.to_int() & self._bitmask(netmask)
 		self.netmask = netmask
-		self.last_ipnum = self.ipnum + self.block_size() - 1
 
 	@staticmethod
 	def parse(cidr: str):
+		"""
+		Parse a string in cidr-notation and return a IPAddressBlock describing the ip-segment
+		If the string is not in cidr-notation a ValueError is raised
+		"""
+		
 		match = re.match("^" + IPAddressBlock.CIDR_REGEXP + "$", cidr)
 		if not match:
 			raise ValueError("%s is no valid cidr-notation" % cidr)
@@ -113,13 +181,22 @@ class IPAddressBlock:
 		return IPAddressBlock(ip, suffix)
 	
 	def block_size(self):
+		"""
+		Return the size of the ip-address-block. E.g. the size of someip/24 is 256
+		"""
 		return 2 ** (32 - self.netmask)
 	
 	def first_address(self):
+		"""
+		Return the first ip-address of the ip-block
+		"""
 		return IPAddress.from_int(self.ipnum)
 
 	def last_address(self):
-		return IPAddress.from_int(self.last_ipnum)
+		"""
+		Return the last ip-address of the ip-block
+		"""
+		return IPAddress.from_int(self.ipnum + self.block_size() - 1)
 
 	def _bitmask(self, netmask):
 		ones = lambda x: (1 << x) - 1
@@ -127,15 +204,27 @@ class IPAddressBlock:
 		return ones(32) ^ ones(32 - netmask)
 	
 	def __repr__(self):
+		"""
+		Conforming to python style-guide, eval(repr(obj)) equals obj
+		"""
 		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 __str__(self):
+		"""
+		Return a string in cidr-notation
+		"""
+		return str(IPAddress.from_int(self.ipnum)) + "/" + str(self.netmask)
 	
 	def __contains__(self, ip):
 		return (ip.to_int() & self._bitmask(self.netmask)) == self.ipnum
 
 class ReservedIPBlocks:
+	"""
+	To avoid magic values and save developers some research this class contains several constants
+	describing special network-segments and some is_-methods to check if an ip is in the specified segment.
+	"""
+	
+ 	# a list of ip-addresses that can be used in private networks
 	PRIVATE_IP_SEGMENTS = [
 		IPAddressBlock.parse(block)
 		for block in

code/ID2TLib/Ports.py

@@ -40,7 +40,24 @@ class PortSelectionStrategy:
 			return random.randrange(port_range.start, port_range.stop)
 
 class PortSelector:
+	"""
+	This class simulates a port-selection-process. Instances keep a list of port-numbers they generated so
+	the same port-number will not be generated again.
+	"""
+	
 	def __init__(self, port_range, select_function):
+		"""
+		Create a PortSelector given a range of ports to choose from and a function that chooses the next port
+		
+		:param port_range: a range-object containing the range of ports to choose from
+		:param select_function: a function that receives the port_range and selects a port
+		"""
+		
+		if len(port_range) == 0:
+			raise ValueError("cannot choose from an empty range")
+		if x.start not in range(1, 65536) or x.stop not in range(1, 65536):
+			raise ValueError("port_range is no subset of the valid port-range")
+		
 		self.port_range = port_range
 		
 		self._select_port = select_function
@@ -48,6 +65,7 @@ class PortSelector:
 		self.generated = []
 	
 	def select_port(self):
+		# do this check to avoid endless loops
 		if len(self.generated) == len(self.port_range):
 			raise RuntimeError("All %i port numbers were already generated, no more can be generated" % len(self.port_range))
 		
@@ -58,25 +76,39 @@ class PortSelector:
 				self.generated.append(port)
 				return port
 	
-	def is_port_in_use(self, port):
+	def is_port_in_use(self, port: int):
 		return port in self.generated
 	
-	def undo_port_use(self, port):
+	def undo_port_use(self, port: int):
 		if port in self.generated:
 			self.generated.remove(port)
 		else:
 			raise ValueError("Port %i is not in use and thus can not be undone" % port)
 	
-	def reduce_size(self, size):
+	def reduce_size(self, size: int):
+		"""
+		Reduce the list of already generated ports to the last <size> generated.
+		If size if bigger than the number of generated ports nothing happens.
+		"""
 		self.generated = self.generated[-size:]
 	
 	def clear(self):
+		"""
+		Clear the list of generated ports. As of now this does not reset the state of the selection-function
+		"""
 		self.generated = []
 	
 	def clone(self):
 		return copy.deepcopy(self)
 
 class ProtocolPortSelector:
+	"""
+	This class contains a method to select ports for udp and tcp. It generally consists of the port-selectors, one
+	for tcp and one for udp. For convenience this class has a __getattr__-method to call methods on both selectors
+	at once. E.g, clear() does not exist for ProtocolPortSelector but it does for PortSelector, therefore
+	protocolPortSelector.clear() will call clear for both port-selectors.
+	"""
+	
 	def __init__(self, port_range, select_tcp, select_udp = None):
 		self.tcp = PortSelector(port_range, select_tcp)
 		self.udp = PortSelector(port_range, select_udp or select_tcp)
@@ -112,18 +144,23 @@ class ProtocolPortSelector:
 	def __getattr__(self, attr):
 		val = getattr(self.tcp, attr)
 		
-		if callable(val):
+		if callable(val): # we proprably got a method here
 			tcp_meth = val
 			udp_meth = getattr(self.udp, attr)
 			
 			def double_method(*args, **kwargs):
 				return (tcp_meth(*args, **kwargs), udp_meth(*args, **kwargs))
 			
-			return double_method
-		else:
+			return double_method # calling this function will call the method for both port-selectors
+		else: # we have found a simple value, return a tuple containing the attribute-value from both port-selectors
 			return (val, getattr(self.udp, attr))
 
 class PortSelectors:
+	"""
+	To save some time this class contains some of the port-selection-strategies found in the wild. It is recommend to use
+	.clone() to get your personal copy, otherwise two parts of your code might select ports on the same port-selector which
+	is something your might want to avoid.
+	"""
 	LINUX = ProtocolPortSelector(PortRanges.LINUX, PortSelectionStrategy.random)
 	APPLE = ProtocolPortSelector(PortRanges.DYNAMIC_PORTS,
 			PortSelectionStrategy.sequential(),