import ipaddress as ipa
import pandas as pd
import json

FULL_VULNERABLE		= 0x3
PART_VULNERABLE		= 0X1
NOT_VULNERABLE		= 0X0


def is_local_address(addr):

	if addr == 'local': return True

	local_networks = [ipa.IPv4Network('192.168.5.0/24')]

	for net in local_networks:
		try:
			if ipa.ip_address(addr) in net: return True
		except:
			break

	return False

def read_net_graph():
	rules = []
	input_rules = []

	with open("/tmp/mud_ovs_rules.txt","r") as file:
		for line in file:
			input_rules.append(line.strip())

	for line in input_rules:

		defs = line.split(',')[1:]
		
		if defs[-1] != 'action=drop':
			defs = defs[:-1]

			rule_line = {}

			for spec in defs:

					aux = spec.split('=')

					if(len(aux) == 2):
						rule_line[aux[0]] = aux[1]
					else:
						rule_line["transport"] = aux[0].lower()

			if rule_line != {}:
				# FILL EMPTY FIELDS

				try:
					rule_line['local_src'] = is_local_address(rule_line['nw_src'])
				except:
					rule_line['local_src'] = True

				try:
					if rule_line['tp_dst'] is None:
						rule_line['tp_dst'] = 'any'
				except:
					rule_line['tp_dst'] = 'any'

				try:
					if rule_line['tp_src'] is None:
						rule_line['tp_src'] = 'any'
				except:
					rule_line['tp_src'] = 'any'

				try:
					if rule_line['nw_src'] is None:
						rule_line['nw_src'] = 'local'
						rule_line['local_src'] = True
				except:
					rule_line['nw_src'] = 'local'
					rule_line['local_src'] = True

				try:
					if rule_line['nw_dst'] is None:
						rule_line['nw_dst'] = 'local'
				except:
					rule_line['nw_dst'] = 'local'

				rules.append(rule_line)

	return pd.DataFrame(rules)


def read_mtds():
	return pd.read_json('/tmp/mtds.json')
	

def compare_transport(malicious, edge):
	edge = str(edge).lower()
	malicious = str(malicious).lower()

	if malicious == 'any' or malicious == '' or edge == malicious:
		return FULL_VULNERABLE
	elif edge == 'any' or edge == '':
		return PART_VULNERABLE
	else:
		return NOT_VULNERABLE

def compare_port(malicious, edge):

		malicious = str(malicious)
		edge = str(edge)

		if malicious == 'any' or pd.isnull(malicious) or malicious == edge:
			return FULL_VULNERABLE

		elif edge == 'any' or pd.isnull(edge):
			return PART_VULNERABLE

		else:
			return NOT_VULNERABLE

def compare_networks(malicious, edge):
	if malicious == 'local' and is_local_address(edge):
		return FULL_VULNERABLE
	else:
		try:
			malicious = ipa.ip_address(malicious)
			edge = ipa.ip_address(edge)

			if malicious == edge:
				return FULL_VULNERABLE
			else:
				return NOT_VULNERABLE
		except:
			return NOT_VULNERABLE

def compare_edges(ace, edges):
	vulnerabilities = []

	for index,edge in edges.iterrows():

		nw_dst = compare_networks(ace['nw_dst'], edge['nw_dst'])
		#print(f"nw_dst: {nw_dst}, ace: {ace['nw_dst']}, edge: {edge['nw_dst']}")

		transport = compare_transport(ace['transport'].lower(), edge['transport'].lower())

		t_src = compare_port(ace['tp_src'], edge['tp_src'])

		t_dst = compare_port(ace['tp_dst'], edge['tp_dst'])

		classification = nw_dst & transport & t_src & t_dst

		if classification:
			vulnerabilities += [{"classification": classification, "mtd": ace, "edge":edge}]

	return vulnerabilities

def address_vulnerabilities():

	net_graph = read_net_graph()
	mtds = read_mtds()
	vulnerabilities = []

	# VERIFY IF MTD SRC ADDRESS IS THE ABSTRACTION LOCAL
	for index,ace in mtds.iterrows():

		edges = []

		# IF TRUE, TEST FOR ALL SRC ADDRESS EDGES
		if ace['nw_src'] == 'local':
			mask_filter = net_graph['local_src'] == True
			edges = net_graph[mask_filter]
			

		# IF FALSE, TEST ONLY BY SPECIFIC SRC ADDRES
		else:
			mask_filter = net_graph['nw_src'] == ace['nw_src']
			edges = net_graph[mask_filter]

		aux = compare_edges(ace, edges)

		if aux:
			vulnerabilities += aux

	return pd.DataFrame(vulnerabilities)

def mount_ovs_rule(nw_src,nw_dst,transport,tp_src,tp_dst):
	rule = f"priority=200"

	if transport is not None and transport != 'any':
		rule += f",{transport}"

	if nw_src is not None and nw_src != 'any':
		rule += f",nw_src={nw_src}"

	if nw_dst is not None and nw_dst != 'any':
		rule += f",nw_dst={nw_dst}"

	if tp_src is not None and tp_src != 'any':
		rule += f",tp_src={tp_src}"

	if tp_dst is not None and tp_dst != 'any':
		rule += f",tp_dst={tp_dst}"

	rule += ",action=drop\n"

	return rule

if __name__ == '__main__':
	a = address_vulnerabilities()

	rules = []

	for index,aux in a.iterrows():
		if aux['classification'] == 3:
			rules.append(mount_ovs_rule(aux['edge']['nw_src'], aux['edge']['nw_dst'],aux['edge']['transport'],aux['edge']['tp_src'],aux['edge']['tp_dst']))
		else:
			rules.append(mount_ovs_rule(aux['edge']['nw_src'], aux['edge']['nw_dst'],aux['mtd']['transport'],aux['mtd']['tp_src'],aux['mtd']['tp_dst']))

	rules = list(dict.fromkeys(rules))

	with open("/tmp/mtd_blocks_ovs.txt",'w') as file:
		file.write(''.join(rules))

	#print(a)
	#a.to_json("/tmp/blocks.json")