#!/usr/bin/env python

import subprocess
import argparse
import configparser
import sys

from pathlib import Path
from string import Template
from typing import Iterable, TypedDict, NamedTuple

TEMPLATE_STR = """[Interface]
PrivateKey = $private_key
Address = $address
DNS = $dns

[Peer]
PublicKey = $remote_public_key
PresharedKey = $psk
AllowedIPs = $allowed_ips
Endpoint = $endpoint"""
TEMPLATE = Template(TEMPLATE_STR)
ENDPOINT = "127.0.0.1:51820"
PUBLIC_KEY = "REPLACE_WITH_SERVER_PUBLIC_KEY"
DNS = "127.0.0.1"


class ConfigArgs(TypedDict):
    remote_public_key: str
    public_key: str
    private_key: str
    psk: str
    address: str
    allowed_ips: str
    endpoint: str
    dns: str


class Config(NamedTuple):
    public_key: str
    psk: str
    config: str


def main():
    # Query for connection info (Public key from server, Address, AllowedIPs, etc...)
    args = parse_args()
    config = load_config(args.config)

    if args.outfile and args.outfile.exists():
        print("The file {args.outfile} already exists. Do you want to overwrite? y/N ", end="")
        overwrite_input = input()
        overwrite = overwrite_input.strip()[0].lower() == 'y'
        if not overwrite:
            sys.exit(0)

    address = args.address
    allowed_ips = args.allowed_ips
    public_key = args.public_key or config["wireguard"]["public_key"]
    endpoint = args.endpoint or config["wireguard"]["endpoint"]
    dns = args.dns or config["wireguard"]["dns"]
    psk = args.psk
    public_key, psk, config = generate_config(
        address,
        allowed_ips,
        public_key,
        endpoint,
        dns,
        psk,
    )
    if not args.quiet:
        print_config(public_key, psk, config)
    if args.outfile:
        save_config_file(args.outfile, config)


def print_config(public_key: str, psk: str, config: str):
    print(f"### Public Key:\n{public_key}", end="\n\n")
    print(f"### PSK:\n{psk}", end="\n\n")
    print(f"### Wireguard Config:\n{config}", end="\n\n")


def generate_config(
    address: str,
    allowed_ips: Iterable[str],
    remote_public_key: str,
    endpoint: str,
    dns: str,
    psk: str | None,
) -> Config:
    allowed_ips = ", ".join(allowed_ips)
    private_key = get_private_key()
    public_key = get_public_key(private_key)
    psk = psk or get_psk()
    config = {
        "private_key": private_key,
        "public_key": public_key,
        "remote_public_key": remote_public_key,
        "psk": psk,
        "address": address,
        "allowed_ips": allowed_ips,
        "endpoint": endpoint,
        "dns": dns,
    }
    config_string = TEMPLATE.substitute(**config)
    return Config(public_key, psk, config_string)


def get_private_key() -> str:
    result = subprocess.run(["wg", "genkey"], stdout=subprocess.PIPE)
    private_key = result.stdout.strip()
    return private_key.decode("utf-8")


def get_public_key(private_key: str) -> str:
    process = subprocess.Popen(
        ["wg", "pubkey"], stdin=subprocess.PIPE, stdout=subprocess.PIPE
    )
    result = process.communicate(input=private_key.encode("utf-8"))
    public_key = result[0].strip()
    return public_key.decode("utf-8")


def get_psk() -> str:
    result = subprocess.run(["wg", "genpsk"], stdout=subprocess.PIPE)
    psk = result.stdout.strip()
    return psk.decode("utf-8")


def save_config_file(filename: Path, config: str):
    filename.touch()
    filename.chmod(0o600)
    filename.write_text(config)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="WireGuard Config File Generator")
    parser.add_argument("address", help="The remote wireguard server hostname or address")
    parser.add_argument("allowed_ips", nargs="+", help="The list of allowed IP ranges for the routed network")
    parser.add_argument("-c", "--config", default="config.ini", type=Path, help="The config file to use for server configuration")
    parser.add_argument("-o", "--outfile", type=Path, help="The file to write to")
    parser.add_argument("-q", "--quiet", default=False, action="store_true", help="Suppress all output")
    parser.add_argument("-p", "--psk", help="The Pre-Shared Key to use")
    parser.add_argument("-e", "--endpoint", help="The Wireguard server to use")
    parser.add_argument("-d", "--dns", help="The DNS server to use")
    parser.add_argument("-P", "--public-key", help="The Servers public key to use")
    return parser.parse_args()


def load_config(filename: Path) -> configparser.ConfigParser:
    config = configparser.ConfigParser()
    if not filename.exists():
        config["wireguard"] = {
            "endpoint": ENDPOINT,
            "public_key": PUBLIC_KEY,
            "dns": DNS,
        }
        with filename.open("w") as fp:
            config.write(fp)

    config.read(filename)
    return config


if __name__ == "__main__":
    main()