# 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
from ipaddress import IPv4Address, IPv4Network
import tomllib
import jinja2

import dns
from config import Config

def load_firewall(config, run):
    run("nft -f templates/rules.nft")

    run(f"nft add element ip filter local_range {{ {config.networks[config.local_network].local_range} }}")
    run(f"nft add element ip filter local_translated_range {{ {config.networks[config.local_network].local_translated_range} }}")
    for net in config.remote_networks:
        run(f"nft add element ip filter remote_range {{ {net} }}")

    for (loc, trans) in config.networks[config.local_network].translation_dict:
        run(f"nft add element ip filter ip_map_snat {{ {loc} : {trans} }}")
        run(f"nft add element ip filter ip_map_dnat {{ {trans} : {loc} }}")


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

    with open("wg-pn.conf", "w") as f:
        f.write(template.render(
            private_key=config.private_key,
            listen_port=config.wg_listen_port,
            wireguard_address=config.local_wireguard_address,
            peers=peers
        ))


def main():
    dry_run = False
    if dry_run:
        run = print
    else:
        run = os.system

    config = Config("/config/config.toml")
    load_firewall(config, run)
    load_wireguard(config, run)

    run("wg-quick up ./wg-pn.conf")
    dns.run(config, port=5353)

if __name__ == "__main__":
    main()