#!/usr/bin/env python3

import re, subprocess


ufw = "/usr/sbin/ufw"
ufw_command = [ufw]
ufw_delete = ufw_command + ["--force", "delete"]

def read_hosts():
	hosts = set()
	pattern = re.compile(r"([0-9.:]+)\s+\w+-(sealer|gateway)")
	with open("/etc/hosts") as file:
		for line in file:
			m = re.match(pattern,line)
			if m:
				hosts.add(m.group(1))
	return hosts

def read_ufw():
	pattern = re.compile(r"\[ *(\d+)\] 30303.+ALLOW IN\s+([0-9.:]+)")
	rules = {}
	ufw = subprocess.run(["/usr/sbin/ufw", "status", "numbered"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
	lines = ufw.stdout.decode('utf-8').split("\n")
	for line in lines:
		m = re.match(pattern,line)
		if m:
			rules[m.group(2)] = m.group(1)
	return rules


def remove_old(allowed, actual):
	to_remove = list(set(actual.keys()) - allowed)
	to_remove.sort(reverse=True, key=actual.get)
	for rule in to_remove:
		subprocess.run(ufw_delete + [actual[rule]], check=True)


def add_allowed(allowed):
	for host in allowed:
		command = ufw_command + ["allow", "from", host, "to", "any", "port", "30303"]
		subprocess.run(command, check=True)

allowed = read_hosts()
actual = read_ufw()
remove_old(allowed,actual)
add_allowed(allowed)