PREPARE/prepare/group.py

import logging
import math

logger = logging.getLogger("pra_framework")


class Group(object):
	"""
	A Group represents a group of IP addresses regarding the PRA.
	It's generally connected to a marker value for identification purposes.
	A Group contains single addreses (group_type == Group.GROUP_TYPE_SINGLE_ADDRESSES)
	OR address ranges (group_type == Group.GROUP_TYPE_CIDR)
	"""
	GROUP_TYPE_SINGLE_ADDRESSES = 0
	GROUP_TYPE_CIDR = 1

	def __init__(self,
				ip_network_object=None,
				ip_host_objects=None,
				response_count=0):
		"""
		Given Parameters: The IPs this group is made of.
		ip_network_bytes -- bytes like b"1234" to create this group
		ip_host_objects -- IPs to be used to create this group
		"""
		self.ip_network = None
		self.ip_hosts = None

		if ip_network_object is not None:
			self.ip_network = ip_network_object
		else:
			self.ip_hosts = ip_host_objects
		# "+1"-groups won't get scanned, the amount of responses is derived implicitly
		self.is_plus1 = False
		# store "+1"-Subgroup for faster access (marker bytes OR group itself)
		self.plus1_subgroup = None
		# logger.debug("group network/addresses: %r/%r" % (self.ip_network, self.ip_hosts))
		# full marker as bytes
		self.marker_bytes = None
		self.marker_value_int = None
		# how many times the group is counted in the report
		self.response_count = response_count
		# needed to create additional blacklists for other groups which do not need
		# to scan this addresses again (already placed in separate group)
		self.top_group = None
		self.subgroups = set()

	def _get_grouptype(self):
		return Group.GROUP_TYPE_CIDR if self.ip_network is not None else Group.GROUP_TYPE_SINGLE_ADDRESSES

	group_type = property(_get_grouptype)

	def _get_amount_addresses(self):
		return self.ip_network.num_addresses if self.group_type == Group.GROUP_TYPE_CIDR else len(self.ip_hosts)

	amount_addresses = property(_get_amount_addresses)

	def add_subgroup(self, group):
		group.top_group = self
		#group.top_group_markervalue_bytes = self.markervalue_bytes
		self.subgroups.add(group)

	def create_subgroups(self, amount_subgroups, ipv4_addresses=[], use_plus1=False):
		"""
		Create amount subgroups based on amount_subgroups and response counts.
		amount_subgroups -- amount of subgroups to create, this includes the +1 group
			e.g. amount_subgroups = 4 = subgroups + "+1"
		ipv4_addresses -- list of single IP addresses as IP objects to be added as dedicated group.
		use_plus1 -- define last subgroup as +1 group
		return -- created subgroups as list
		"""
		subgroups = []

		if self.response_count >= self.amount_addresses:
			if self.response_count > self.amount_addresses:
				logger.warning("got more responses than amount addresses, too much noise?"
								" responses/# addresses = %d/%d" % (self.response_count, self.amount_addresses))
			return subgroups
		# check if this is a single address group (CIDR prefix is 32 or just 1 single address)
		if self.amount_addresses <= 1:
			return subgroups

		if self.group_type == Group.GROUP_TYPE_SINGLE_ADDRESSES:
			single_addr_amount = len(self.ip_hosts)
			"""
			logger.debug("subgrouping single addresses: addresses/target subgroups = %d/%d" %
						(single_addr_amount, amount_subgroups))
			"""
			if single_addr_amount == 0:
				logger.warning("no host for subgrouping!")

			# split by amount of adresses
			groupsize_single = math.floor(single_addr_amount / amount_subgroups)
			# at minimum 1 group
			groupsize_single = max(1, groupsize_single)

			for i in range(0, single_addr_amount, groupsize_single):
				# not enough room for more groups, add all remaining addresses
				if len(subgroups) + 1 >= amount_subgroups:
					subgroup = Group(ip_host_objects=self.ip_hosts[i:])
					subgroups.append(subgroup)
					break
				else:
					sub_ips = self.ip_hosts[i: i + groupsize_single]
					# this should not happen
					if len(sub_ips) != 0:
						subgroup = Group(ip_host_objects=sub_ips)
						subgroups.append(subgroup)
					else:
						break

		elif self.group_type == Group.GROUP_TYPE_CIDR:
			subgroups_single_addresses = 0
			#logger.debug("picking IPv4 addresses which belong to this groups (searching through: %d)" % len(ip_host_objects))
			single_addresses_for_regrouping = ipv4_addresses
			# amount of single addresses to be re-groupted
			single_addr_amount = len(single_addresses_for_regrouping)

			if single_addr_amount > 0:
				if len(single_addresses_for_regrouping) > 0:
					#logger.debug("re-grouping %d single addresses for nw %r" % (single_addr_amount, self.ip_network))
					#for single_addr in single_addresses_for_regrouping:
					#	logger.debug("%r" % single_addr)
					# calculate amount of subgroups for single addresses:
					# (amount of single addreses / total group address space) * amount of groups
					subgroups_single_addresses = math.floor((single_addr_amount / self.ip_network.num_addresses) *
													amount_subgroups)
					# at minimum 1 group
					subgroups_single_addresses = max(1, subgroups_single_addresses)
					# not more groups than single addresses
					subgroups_single_addresses = min(single_addr_amount, subgroups_single_addresses)
					groupsize_single = math.floor(single_addr_amount / subgroups_single_addresses)
					"""
					logger.debug("creating single addresses groups,"
								"addr total=%d, groups total=%d, group size=: %d/%d/%d" %
								(single_addr_amount,
								subgroups_single_addresses,
								groupsize_single))
					"""
					for i in range(0, single_addr_amount, groupsize_single):
						if len(subgroups) + 1 >= subgroups_single_addresses:
							group = Group(ip_host_objects=single_addresses_for_regrouping[i:])
							#logger.debug("adding single addresses group: %r" % group)
							subgroups.append(group)
							break
						else:
							addresses = single_addresses_for_regrouping[i: i + groupsize_single]

							if len(addresses) != 0:
								group = Group(ip_host_objects=addresses)
								#logger.debug("adding single addresses group: %r" % group)
								subgroups.append(group)
							else:
								# no more groups to split up
								break

			# round to next lower integer with 2**x <= (amount of groups for cidr)
			# Example: 16/5 (ip/groups) = 5 -> 2**_log_2(5)_ = 4 top groups [4,4,4,4]
			# (which are splittet later on) -> [4,4,4,2,2]
			cidr_bits_plus = math.floor(math.log(
									amount_subgroups - subgroups_single_addresses,
									2))
			# not more CIDR bits than available
			cidr_bits_plus = min(32 - self.ip_network.prefixlen, cidr_bits_plus)
			"""
			logger.debug("current prefix=%d, CIDR bits plus=%d, amount subgroups=%d, single addresses=%d" %
						(self.ip_network.prefixlen, cidr_bits_plus, amount_subgroups, subgroups_single_addresses))
			"""
			# create subnets: e.g. 1.2.3.0/24 -> CIDR+1 -> 1.2.3.0/25, 1.2.3.128/25
			cidr_nw = self.ip_network.subnets(prefixlen_diff=cidr_bits_plus)
			cidr_nw_len_at_start = len(cidr_nw)
			# amount of times we reached CIDR /32
			splitfail_cnt = 0
			# logger.debug("re-splitting groups for CIDR, initial groupsize/target amount of groups/CIDR bits +x: %d/%d/%d" %
			#			(cidr_nw_len_at_start, amount_groups_for_cidr, cidr_bits_plus))
			subgroup_len = len(subgroups)

			while len(cidr_nw) + subgroup_len < amount_subgroups:
				# split subgroups until we have enough of them
				# [A,B,C,D] -> split A by 1 bit -> [B,C,D,a,a] -> split B by 1 bit -> [C,D,a,a,b,b] ...
				group_to_split = cidr_nw[0]
				del cidr_nw[0]

				if group_to_split.prefixlen == 32:
					# nothing to split: re-append to the end
					cidr_nw.append(group_to_split)
					splitfail_cnt += 1
					#logger.debug("can't create subnets anymore: /32 reached for %r" % group_to_split)
					if splitfail_cnt > len(cidr_nw):
						# logger.warning("too many split fails: single addresses reached?")
						break
				else:
					subnets = group_to_split.subnets(prefixlen_diff=1)
					if subgroup_len + len(cidr_nw) + len(subnets) > amount_subgroups:
						logger.debug("!!! stopping CIDR subgrouping: split would increase max number")
						cidr_nw.append(group_to_split)
						break
					#logger.debug("new subgroups: %d" % len(subnets))
					# append subnet to the end and try next
					cidr_nw.extend(subnets)

			if cidr_nw_len_at_start == len(cidr_nw):
				# logger.debug("no CIDR groups have been re-splitted (perhaps there were enough (2**x)): %d" %
				#			cidr_nw_len_at_start)
				pass

			# then add CIDR based groups
			for nw in cidr_nw:
				#logger.debug("g5: adding sub: %r" % nw)
				group = Group(ip_network_object=nw)
				subgroups.append(group)

		if use_plus1 and len(subgroups) > 1:
			subgroups[-1].is_plus1 = True
		"""
		if len(subgroups) < 600:
			for g in subgroups:
				logger.debug("%r" % g)
		"""
		# state should now be: [single address groups]+[CIDR groups]
		return subgroups

	def _addr_get(self):
		addr = []
		if self.ip_network is not None:
			addr.append(self.ip_network.compressed)
		else:
			addr.extend([host.compressed for host in self.ip_hosts])
		return addr
	# return all addresses as list of strings (mixed CIDR and single notation)
	addresses = property(_addr_get)

	def _addr_get_single_bytes(self):
		if self.group_type == Group.GROUP_TYPE_CIDR:
			return self.ip_network.hosts
		else:
			return [host.packed for host in self.ip_hosts]
	# return all addresses als list of bytes (only single non-CIDR notation)
	addresses_single_bytes = property(_addr_get_single_bytes)

	def __repr__(self):
		if self.group_type == Group.GROUP_TYPE_CIDR:
			# "1.2.3.4/x"
			return self.ip_network.compressed
		else:
			# "1.2.3.4 1.2.3.5 ..."
			return " ".join([host.compressed for host in self.ip_hosts])