#!/usr/bin/python3
"""Minimal root helper for VPN Now's WireGuard lifecycle."""

import os
import base64
import binascii
import ipaddress
import re
import stat
import subprocess
import sys

STATE_DIR = "/run/vpnnow"
ROOT_CONFIG = STATE_DIR + "/vpnnow.conf"
WG_QUICK = "/usr/bin/wg-quick"
ALLOWED = {
    "Interface": {"PrivateKey", "Address", "DNS", "MTU", "ListenPort"},
    "Peer": {"PublicKey", "PresharedKey", "AllowedIPs", "Endpoint", "PersistentKeepalive"},
}


def validate_value(key, value):
    value = value.strip()
    if not value or any(char in value for char in (";", "`", "$", "\\", "\"", "'")):
        fail(f"Invalid value for WireGuard option: {key}")
    try:
        if key in {"PrivateKey", "PublicKey", "PresharedKey"}:
            if len(base64.b64decode(value, validate=True)) != 32:
                raise ValueError
        elif key == "Address":
            for item in value.split(","):
                ipaddress.ip_interface(item.strip())
        elif key == "DNS":
            for item in value.split(","):
                ipaddress.ip_address(item.strip())
        elif key == "AllowedIPs":
            for item in value.split(","):
                ipaddress.ip_network(item.strip(), strict=False)
        elif key == "MTU":
            if not 576 <= int(value) <= 9000:
                raise ValueError
        elif key in {"ListenPort", "PersistentKeepalive"}:
            if not 0 <= int(value) <= 65535:
                raise ValueError
        elif key == "Endpoint":
            match = re.fullmatch(r"(?:\[[0-9A-Fa-f:]+\]|[A-Za-z0-9.-]+):(\d{1,5})", value)
            if not match or not 1 <= int(match.group(1)) <= 65535:
                raise ValueError
    except (ValueError, TypeError, binascii.Error):
        fail(f"Invalid value for WireGuard option: {key}")


def fail(message):
    print(message, file=sys.stderr)
    raise SystemExit(1)


def validate_source(path, caller_uid):
    expected = f"/run/user/{caller_uid}/vpn.now/vpnnow.conf"
    if os.path.realpath(path) != expected:
        fail("Refusing a configuration outside the caller's runtime directory.")
    info = os.lstat(path)
    if not stat.S_ISREG(info.st_mode) or info.st_uid != caller_uid or info.st_mode & 0o077:
        fail("The VPN configuration must be a caller-owned 0600 regular file.")

    section = None
    seen = set()
    with open(path, "r", encoding="utf-8") as handle:
        content = handle.read(65537)
    if len(content.encode("utf-8")) > 65536:
        fail("The VPN configuration is too large.")
    for raw in content.splitlines():
        line = raw.strip()
        if not line or line.startswith("#"):
            continue
        match = re.fullmatch(r"\[(Interface|Peer)\]", line)
        if match:
            section = match.group(1)
            seen.add(section)
            continue
        if section is None or "=" not in line:
            fail("The VPN configuration has invalid syntax.")
        key, value = (part.strip() for part in line.split("=", 1))
        if key not in ALLOWED[section]:
            fail(f"Unsupported WireGuard option: {key}")
        validate_value(key, value)
    if seen != {"Interface", "Peer"}:
        fail("The VPN configuration is incomplete.")
    return content


def write_root_config(content):
    os.makedirs(STATE_DIR, mode=0o700, exist_ok=True)
    fd = os.open(ROOT_CONFIG, os.O_WRONLY | os.O_CREAT | os.O_TRUNC | os.O_NOFOLLOW, 0o600)
    with os.fdopen(fd, "w", encoding="utf-8") as handle:
        handle.write(content)
    os.chown(ROOT_CONFIG, 0, 0)
    os.chmod(ROOT_CONFIG, 0o600)


def erase_root_config():
    try:
        size = os.path.getsize(ROOT_CONFIG)
        with open(ROOT_CONFIG, "r+b", buffering=0) as handle:
            handle.write(b"\0" * size)
        os.unlink(ROOT_CONFIG)
    except FileNotFoundError:
        pass


def run(action):
    if not os.path.isfile(WG_QUICK):
        fail("wireguard-tools is not installed.")
    result = subprocess.run(
        [WG_QUICK, action, ROOT_CONFIG],
        env={"PATH": "/usr/bin:/bin:/usr/sbin:/sbin", "LANG": "C", "HOME": "/root"},
        text=True,
    )
    if result.returncode:
        raise SystemExit(result.returncode)


def down():
    if os.path.exists(ROOT_CONFIG):
        try:
            if os.path.exists("/sys/class/net/vpnnow"):
                run("down")
        finally:
            erase_root_config()


def main():
    if os.geteuid() != 0:
        fail("This helper must be started by polkit.")
    if len(sys.argv) < 2:
        fail("Usage: vpnnow-tunnel-helper up CONFIG | down | cleanup")
    operation = sys.argv[1]
    if operation == "cleanup" and "PKEXEC_UID" not in os.environ:
        down()
        return
    try:
        caller_uid = int(os.environ["PKEXEC_UID"])
    except (KeyError, ValueError):
        fail("Missing polkit caller identity.")
    if caller_uid <= 0:
        fail("Invalid polkit caller identity.")

    if operation == "up" and len(sys.argv) == 3:
        content = validate_source(sys.argv[2], caller_uid)
        down()
        write_root_config(content)
        try:
            run("up")
        except BaseException:
            erase_root_config()
            raise
    elif operation == "down" and len(sys.argv) == 2:
        down()
    else:
        fail("Unsupported tunnel operation.")


if __name__ == "__main__":
    main()
