#!/usr/bin/python3

import json
import subprocess
import sys
from ipaddress import ip_address, ip_network
from itertools import islice
from pathlib import Path

wg_dir  = Path('/etc/wireguard')
wg_json = wg_dir / 'wg0.json'
wg_conf = wg_dir / 'wg0.conf'

iface = 'enp3s0'

ipv4_prefix = '/24'
ipv6_prefix = '/64'

def genkey():
    return subprocess.check_output(['wg', 'genkey']).strip().decode('utf-8')

def pubkey(key):
    return subprocess.run(['wg', 'pubkey'], input=key, encoding='utf-8', capture_output=True).stdout.strip()

def ipv4(cfg):
    taken = [ip_address(cfg['ipv4'])] + [ip_address(c['ipv4']) for c in cfg['clients']] if cfg else []
    for ip in ip_network('192.168.0.0' + ipv4_prefix).hosts():
        if ip not in taken:
            return str(ip)
    raise Exception('no ipv4 left')

def ipv6(cfg):
    taken = [ip_address(cfg['ipv6'])] + [ip_address(c['ipv6']) for c in cfg['clients']] if cfg else []
    for ip in ip_network('fd32:76a6:ec61:577a::0' + ipv6_prefix).hosts():
        if ip not in taken:
            return str(ip)
    raise Exception('no ipv6 left')

def clientconf(cfg, key):
    c = cfg['clients'][-1]
    return (
        f'[Interface]\n'
        f'Address = {c["ipv4"] + ipv4_prefix}\n'
        f'Address = {c["ipv6"] + ipv6_prefix}\n'
        f'DNS = 8.8.8.8, 8.8.4.4, 2001:4860:4860::8888, 2001:4860:4860::8844\n' # Google DNS servers
        f'PrivateKey = {key}\n'
        f'\n'
        f'[Peer]\n'
        f'Endpoint = matthewtran.com:51820\n'
        f'AllowedIPs = 0.0.0.0/0, ::/0\n'
        f'PublicKey = {pubkey(cfg["key"])}\n'
        f'PersistentKeepalive = 15\n'
    )

def serverconf(cfg):
    conf = (
        f'[Interface]\n'
        f'Address = {cfg["ipv4"] + ipv4_prefix}\n'
        f'Address = {cfg["ipv6"] + ipv6_prefix}\n'
        f'ListenPort = 51820\n'
        f'PrivateKey = {cfg["key"]}\n'
        f'PostUp = sysctl -w net.ipv4.ip_forward=1\n'
        f'PostUp = sysctl -w net.ipv6.conf.all.forwarding=1\n'
        f'PostUp = iptables -A FORWARD -i wg0 -j ACCEPT\n'
        f'PostUp = iptables -t nat -A POSTROUTING -o {iface} -j MASQUERADE\n'
        f'PostUp = ip6tables -A FORWARD -i wg0 -j ACCEPT\n'
        f'PostUp = ip6tables -t nat -A POSTROUTING -o {iface} -j MASQUERADE\n'
        f'PostUp = ufw reload\n'
        f'PostDown = sysctl -w net.ipv4.ip_forward=0\n'
        f'PostDown = sysctl -w net.ipv6.conf.all.forwarding=0\n'
        f'PostDown = iptables -D FORWARD -i wg0 -j ACCEPT\n'
        f'PostDown = iptables -t nat -D POSTROUTING -o {iface} -j MASQUERADE\n'
        f'PostDown = ip6tables -D FORWARD -i wg0 -j ACCEPT\n'
        f'PostDown = ip6tables -t nat -D POSTROUTING -o {iface} -j MASQUERADE\n'
        f'PostDown = ufw reload\n'
        f'\n'
    )
    for c in cfg['clients']:
        conf += (
            f'[Peer]\n'
            f'AllowedIPs = {c["ipv4"] + "/32"}\n'
            f'AllowedIPs = {c["ipv6"] + "/128"}\n'
            f'PublicKey = {c["pubkey"]}\n'
            f'\n'
        )
    return conf

if __name__ == '__main__':
    # create initial config if doesn't exist
    if not wg_json.is_file():
        with wg_json.open('w') as file:
            json.dump({
                'ipv4': ipv4(None),
                'ipv6': ipv6(None),
                'key' : genkey(),
                'clients': []
            }, file, indent=4)
            file.write('\n')

    # read config
    with wg_json.open('r') as file:
        cfg = json.load(file)

    # add additional clients
    for c in sys.argv[1:]:
        key = genkey()
        cfg['clients'].append({
            'ipv4'  : ipv4(cfg),
            'ipv6'  : ipv6(cfg),
            'pubkey': pubkey(key),
        })

        with open(c + '.conf', 'w') as file:
            file.write(clientconf(cfg, key)) # qrencode -t ansiutf8 < <conf>

    # generate files
    with wg_json.open('w') as file:
        json.dump(cfg, file, indent=4)
        file.write('\n')

    with wg_conf.open('w') as file:
        file.write(serverconf(cfg))

    # reload new configs
    subprocess.check_output(['systemctl', 'reload', 'wg-quick@wg0.service'])