#!/usr/bin/env nix-shell
#!nix-shell -i python3 -p bunpen -p iproute2 -p networkmanager-split.nmcli -p python3 -p sane-scripts.ip-check -p systemd
# vim: set filetype=python :
"""
sane-vpn: tool to route all system internet traffic through some VPN, or just one application's
and, thanks to cap_net_admin, we can do all this without superuser!
however, systemd --user has poor support for capabilities, hence this here is a bespoke script instead of a service

usage:
sane-vpn [flags ...] operation [operation flags]

flags:
  --verbose
  --no-proxy-dns

operations:
sane-vpn up [REGION]
sane-vpn down [REGION]
sane-vpn do [REGION [-- COMMAND ...] ]
sane-vpn do - [-- COMMAND ...]
sane-vpn --help

special regions:
- none: allow ALL network devices; don't actually re-route traffic through a VPN. useful to make only DNS changes
- unmetered: route traffic PLAINTEXT through whichever unmetered connections (i.e. eth/wifi) are up

idioms:
- sane-vpn do none -- [COMMAND ...]
  - run the command with a stub resolver instead of my recursive resolver, but no VPN.
- sane-vpn up none
  - patch the entire system to use a stub resolver, but no VPN.
- sane-vpn --no-proxy-dns up - -- [COMMAND ...]
  - patch the system to route all traffic over the VPN, but use our stub resolver (still through the VPN) instead of delegating to the VPN owner's resolver
- sane-vpn do unmetered -- [COMMAND ...]
  - run the command using only those interfaces which are unmetered (i.e. ethernet/wifi)
"""

from dataclasses import dataclass
from enum import Enum

import argparse
import json
import logging
import os
import subprocess
import sys

logger = logging.getLogger(__name__)

NM_RESOLVCONF = "/run/NetworkManager/resolv.conf"
RESOLVCONF = "/etc/resolv.conf"

SHELL = os.environ.get("SHELL", "/bin/sh")
HOME = os.environ.get("SUDO_HOME", os.environ.get("HOME", "/root"))
VPN_CONFIG_DIR = os.path.join(HOME, ".config/sane-vpn")
VPN_DIR = os.path.join(VPN_CONFIG_DIR, "vpns")
VPN_DEFAULT_FILE = os.path.join(VPN_CONFIG_DIR, "default")

class ToggleDir(Enum):
    Up = "up"
    Down = "down"

@dataclass
class VpnConfig:
    addr_v4: str|None = None
    dev: str|None = None
    dns: list[str] = ()
    fwmark: int|None = None
    id_: int|None = None
    priority_fwmark: int|None = None
    priority_main: int|None = None

def load_vpns() -> dict[str, VpnConfig]:
    """
    load all known VPN configs
    """

    # e.g.
    # id=3
    # fwmark=10003
    # priorityMain=103
    # priorityFwMark=303
    # addrV4=172.26.55.21
    # name=ovpnd-ukr
    # dns=(46.227.67.134 192.165.9.158)
    vpns = {}
    logger.debug(f"loading VPNs from disk at {VPN_DIR}")
    for fname in os.listdir(VPN_DIR):
        config = VpnConfig()
        for line in open(os.path.join(VPN_DIR, fname), "r").readlines():
            line = line.rstrip()  # remove trailing newline
            sep = line.index("=")
            fieldname = line[:sep]
            value = line[sep+1:]
            if fieldname == "addrV4":
                config.addr_v4 = value
            elif fieldname == "dns":
                config.dns = value.rstrip(")").lstrip("(").split(" ")
            elif fieldname == "fwmark":
                config.fwmark = int(value)
            elif fieldname == "id":
                config.id_ = value
            elif fieldname == "name":
                config.dev = value
            elif fieldname == "priorityFwMark":
                config.priority_fwmark = int(value)
            elif fieldname == "priorityMain":
                config.priority_main = int(value)
            else:
                assert False, f"unrecognized vpn field {fieldname}"
        vpns[fname] = config

    return vpns

def get_default_dev() -> str:
    """
    return the device name over which packets are sent by default
    """
    default_route = subprocess.check_output([
      "ip", "-j", "route", "get", "255"
    ]).decode("utf-8")
    logger.debug(f"default route: {default_route}")
    # looks like [{"dst":"255.0.0.0","gateway":"10.78.79.1","dev":"eth0","prefsrc":"10.78.79.52","flags":[],"uid":1000,"cache":[]}]

    for r in json.loads(default_route):
        return r["dev"]

def nameservers_from_resolvconf(resolv_conf_path: str) -> list[str]:
    """
    return a list of DNS resolvers we could use, i.e. the list of resolvers which were issued to us over DHCP.
    """
    # file looks like:
    # # Generated by NetworkManager
    # search my.domain
    # nameserver a:b:c::e
    # nameserver 1:2:3::4
    # nameserver 192.168.0.1
    # # NOTE: the libc resolver may not support more than 3 nameservers.
    # # The nameservers listed below may not be recognized.
    # nameserver 192.168.0.2

    nameservers = []
    for line in open(resolv_conf_path, "r").readlines():
        line = line.rstrip()  #< remove trailing '\n'
        if line.startswith("nameserver "):
            nameservers.append(line[len("nameserver "):].strip())

    return nameservers

def get_dns_resolvers_for_dev(dev: str) -> list[str]:
    """
    passt/pasta can't proxy to e.g. 127.0.0.53, but it can to 127.0.0.1 and anything routable by the device
    """

    # output looks like IP4.DNS[1]:192.168.0.1
    nmcli_stdout = subprocess.check_output([
        "nmcli", "-t", "-f", "IP4.DNS", "device", "show", dev
    ]).decode("utf-8")
    nameservers = []
    for line in nmcli_stdout.splitlines():
        line = line.strip()
        if not line:
            continue
        _fname, ns = line.split(":")
        nameservers.append(ns)

    if nameservers == []:
        nameservers = nameservers_from_resolvconf(RESOLVCONF)
        if nameservers == []:
            nameservers = [ "127.0.0.1 "]  #< libc default if unspecified

    # old logic from when i was using systemd-resolved (127.0.0.53, which pasta couldn't reach). safe to remove?
    # if "127.0.0.1" in nameservers:
    #     return [ "127.0.0.1" ]

    return nameservers

def get_addr_v4_for_dev(dev: str) -> str | None:
    # output looks like IP4.ADDRESS[1]:192.168.1.72
    nmcli_stdout = subprocess.check_output([
        "nmcli", "-t", "-f", "IP4.ADDRESS", "device", "show", dev
    ]).decode("utf-8")
    addrs = []
    for line in nmcli_stdout.splitlines():
        line = line.strip()
        if not line:
            continue
        _fname, addr = line.split(":")
        if not addr:
            continue
        if "/" in addr:
            addr, _subnet = addr.split("/")
        if not addr:
            continue
        addrs.append(addr)

    return addrs[0] if len(addrs) else None

def get_connected_networks() -> list[str]:
    # output looks like:
    # 
    # eth0:ethernet:connected
    # lo:loopback:connected (externally)
    # ovpnd-us:wireguard:connected (externally)
    # wg-home:wireguard:connected (externally)
    # wlan0:wifi:unmanaged
    ok_networks = []
    nmcli_stdout = subprocess.check_output(["nmcli", "--get-values", "DEVICE,TYPE,STATE", "device",]).decode("utf-8")
    logger.debug(f"nmcli output: {nmcli_stdout}")
    for line in nmcli_stdout.splitlines():
        line = line.strip()
        if not line:
            continue

        device, type_, state = line.split(":")
        logger.debug(f"considering {device} ({type_}, {state})")
        if state != "connected":
            logger.debug(f"skipping {device}: not connected")
            continue
        if type_ == "ethernet":
            # ethernet gets high precedence
            ok_networks = [ device ] + ok_networks
        elif type_ == "wifi":
            # wifi gets low precedence
            ok_networks += [ device ]
        else:
            logger.debug(f"skipping {device}: not ethernet/wifi")

    return ok_networks

def get_vpn(vpn_name: str) -> VpnConfig:
    if vpn_name in ( "default", "", "-", "--" ):
        vpn_name = open(VPN_DEFAULT_FILE, "r").read().strip()
        logger.debug(f"loading default VPN: {vpn_name}")
        return load_vpns()[vpn_name]
    elif vpn_name == "none":
        dev = get_default_dev()
        addr_v4 = get_addr_v4_for_dev(dev)
        dns = get_dns_resolvers_for_dev(dev)
        return VpnConfig(dev=dev or "all", addr_v4=addr_v4, dns=dns)
    elif vpn_name == "unmetered":
        for dev in get_connected_networks():
            dns = get_dns_resolvers_for_dev(dev)
            return VpnConfig(dev=dev, dns=dns)
    else:
        vpns = load_vpns()
        v = vpns.get(vpn_name) or vpns.get(f"ovpnd-{vpn_name}")
        if v:
            return v

    assert False, f"failed to get vpn {vpn_name}"

def vpn_do(config: VpnConfig, cmd: list[str]) -> None:
    net_flags = []
    net_flags += [ "--bunpen-net-dev", config.dev ]
    for d in config.dns:
        net_flags += [ "--bunpen-dns", d ]
    if config.addr_v4:
        net_flags += [ "--bunpen-net-gateway", config.addr_v4 ]

    wrapped_cmd = [
        "bunpen",
        "--bunpen-keep-ipc",
        "--bunpen-keep-pid",
        "--bunpen-try-keep-users",
        "--bunpen-path", "/",
    ] + net_flags + cmd

    logger.debug(f"exec {wrapped_cmd}")
    # exec instead of subprocess so that (hopefully) the process inherits stdin/stdout.
    os.execvp("bunpen", wrapped_cmd)

def vpn_toggle(config: VpnConfig, dir_: ToggleDir, do_ip_check: bool = True) -> None:
    if do_ip_check:
        print("old IP address ...")
        _ = subprocess.run(["sane-ip-check", "--no-upnp", "--retry-duration", "2"])

    if config.priority_main and config.priority_fwmark:
        verb = {
            ToggleDir.Up: "add",
            ToggleDir.Down: "del",
        }[dir_]
        # first, allow all non-default routes (prefix-length != 0) a chance to route the packet.
        # - this allows the wireguard tunnel itself to pass traffic via our LAN gateway.
        # - incidentally, it allows traffic to LAN devices and other machine-local or virtual networks.
        subprocess.check_call([
            "ip", "rule", verb,
            "from", "all",
            "lookup", "main",
            "suppress_prefixlength", "0",
            "priority", str(config.priority_main),
        ])

        # if packet hasn't gone through the wg device yet (fwmark), then move it to the table which will cause it to.
        subprocess.check_call([
            "ip", "rule", verb,
            "not", "from", "all",
            "fwmark", str(config.fwmark),
            "lookup", str(config.id_),
            "priority", str(config.priority_fwmark),
        ])

    dns_toggle(config.dns, dir_)

    if do_ip_check:
        print("new IP address ...")
        subprocess.check_call(["sane-ip-check", "--no-upnp"])

def dns_toggle_bind(dns: list[str], dir_: ToggleDir) -> None:
    if dir_ == ToggleDir.Up:
        formatted_nameservers = "\n".join(f"{ns};" for ns in dns)
        text = f'''
zone . {{
  //v XXX(2025-01-30): BIND9 doesn't allow dnssec-validation per-zone. put this in toplevel `options` instead.
  // dnssec-validation no;  // compatibility: many network-specific DNS servers fail DNSSEC, by design
  type forward;
  forward first;
  forwarders {{
    {formatted_nameservers}
  }};
}};
'''
    elif dir_ == ToggleDir.Down:
        text = ""

    if not os.path.isdir("/run/named/dhcp-configs"):
        logger.info("not restarting bind because it appears to not be in use")
        return

    with open("/run/named/dhcp-configs/210-sane-vpn.conf", "w") as f:
        f.write(text)
    subprocess.check_call([ "systemctl", "restart", "bind" ])

def dns_toggle_hickory(dns: list[str], dir_: ToggleDir) -> None:
    if dir_ == ToggleDir.Up:
        formatted_nameservers = ",\n".join(
            '{ socket_addr = "{ns}:53", protocol = "udp", trust_nx_responses = false }'.replace("{ns}", ns)
            for ns in dns
        )
        text = f'''
[[zones]]
zone = "."
zone_type = "Forward"
stores = {{ type = "forward", name_servers = [
    {formatted_nameservers}
]}}
'''
    elif dir_ == ToggleDir.Down:
        text = ""

    if not os.path.isdir("/var/lib/hickory-dns"):
        logger.info("not restarting hickory-dns because it appears to not be in use")
        return

    with open("/var/lib/hickory-dns/dhcp-configs/sane-vpn.toml", "w") as f:
        f.write(text)
    subprocess.check_call([ "systemctl", "restart", "hickory-dns-localhost" ])

def dns_toggle(dns: list[str], dir_: ToggleDir) -> None:
    dns_toggle_bind(dns, dir_)
    dns_toggle_hickory(dns, dir_)

def main():
    logging.basicConfig()
    logging.getLogger().setLevel(logging.INFO)

    parser = argparse.ArgumentParser(usage=__doc__)
    parser.add_argument("--verbose", action="store_true")
    parser.add_argument("--no-proxy-dns", action="store_true")
    subparsers = parser.add_subparsers(help="what to do")

    up_parser = subparsers.add_parser("up")
    up_parser.set_defaults(subcommand="up")
    up_parser.add_argument("region", nargs="?", default="default", help="which vpn to connect to")

    down_parser = subparsers.add_parser("down")
    down_parser.set_defaults(subcommand="down")
    down_parser.add_argument("region", nargs="?", default="default", help="which vpn to connect to")

    do_parser = subparsers.add_parser("do")
    do_parser.set_defaults(subcommand="do")
    do_parser.add_argument("region", nargs="?", default="default", help="which vpn to connect to")
    do_parser.add_argument("command", nargs="*", default=[ SHELL ], help="command to run inside the VPN")

    args = parser.parse_args()

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    subcommand = getattr(args, "subcommand", None)
    region = getattr(args, "region", "default")

    logger.debug(f"sane-vpn {subcommand} region={region} ...")

    do_ip_check = region != "none"
    vpn_config = get_vpn(region)

    if args.no_proxy_dns:
        vpn_config.dns = []

    if subcommand == "do":
        vpn_do(vpn_config, args.command)
    elif subcommand == "down":
        vpn_toggle(vpn_config, ToggleDir.Down, do_ip_check=do_ip_check)
    elif subcommand == "up":
        vpn_toggle(vpn_config, ToggleDir.Up, do_ip_check=do_ip_check)
    else:
        print(parser.usage)
        sys.exit(1)


if __name__ == '__main__':
    main()
