# PolyculeNetwork
# Copyright (C) 2024 PolyculeConnect
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import os
import ipaddress
import tomllib
import jinja2


dry_run = True
if dry_run:
    run = print
else:
    run = os.system

def load_config(path):
    with open(path, "rb") as f:
        data = tomllib.load(f)
        return data

def load_firewall():
    data = load_config("config.toml")

    run("sudo 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 = str(ipaddress.IPv4Network(data["network"][local_network]["local_range"]))
    local_translated_range = str(ipaddress.IPv4Network(data["network"][local_network]["local_translated_range"]))
    remote_ranges = [str(ipaddress.IPv4Network(data["network"][net]["local_translated_range"])) for net in remote_networks]

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

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


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

    peers = []


    data = load_config("config.toml")
    networks = data["network"].keys()
    local_network = data["local"]
    remote_networks = list(filter(lambda k: k != local_network, networks))
    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)


    with open("wg-pn.conf", "w") as f:
        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
        ))

load_firewall()