import os
import ipaddress
import tomllib
import jinja2

f = open("config.toml", "rb")
data = tomllib.load(f)
f.close()

run = print
# run = os.system

run("nft -f templates/rules.nft")

networks = data["network"].keys()

local_network = data["local"]
remote_networks = list(filter(lambda k: k != local_network, networks))

local_range = data["network"][local_network]["local_range"]
local_translated_range = data["network"][local_network]["local_translated_range"]
remote_ranges = [data["network"][net]["local_translated_range"] for net in remote_networks]

run(f"nft add element ip filter local_range {{ {local_range} }}")
for net in remote_ranges:
    run(f"nft add element ip filter remote_range {{ {net} }}")
run(f"nft add element ip filter local_translated_range {{ {local_translated_range} }}")

for (loc, trans) in zip(ipaddress.IPv4Network(local_range), ipaddress.IPv4Network(local_translated_range)):
    run(f"nft add element ip filter ip_map_dnat {{ {loc} : {trans} }}")
    run(f"nft add element ip filter ip_map_dnat {{ {trans} : {loc} }}")


f = open("templates/wg-pn.conf.j2", "r")
env = jinja2.Environment()
template = env.from_string(f.read())

peers = []

for net in remote_networks:
    peer = {
        "public_key": data["network"][net]["public_key"],
    }

    endpoint = data["network"][net].get("endpoint", "")
    if endpoint != "":
        peer["endpoint"] = endpoint
    

    peer["allowed_ips"] = data["network"][net]["local_translated_range"]
    untranslated_networks = data["network"][net].get("untranslated_networks", "")
    if untranslated_networks != "":
        peer["allowed_ips"] += ", " + untranslated_networks


    peers.append(peer)


f = open("wg-pn.conf", "w")
f.write(template.render(
    private_key=data["network"][local_network]["private_key"],
    listen_port=data["network"][local_network]["listen_port"],
    wireguard_address=data["network"][local_network]["wireguard_address"],
    peers=peers
))
f.close()