"""
ARIA VPN-Satellit — ein WireGuard-Tunnel + ein ganz normaler Satellit dahinter.

Laeuft als eigener Container pro VPN-Verbindung (angelegt von der Diagnostic,
Einstellungen → VPN → WireGuard). Eigener Container = eigener Netz-Namensraum:
  - ueberlappende Kundennetze (jeder hat 192.168.178.0/24) stoeren sich nicht
  - lokale Satelliten (network_mode: host) sehen diese Netze nie

Ablauf:
  1. /vpn/wg.conf lesen (Standard-WireGuard-Datei, wie man sie bekommt)
  2. wg0 selbst aufsetzen (kein wg-quick: DNS/PostUp aus fremden Dateien
     werden ignoriert, und der Weg zum RVS darf NIE durch den Tunnel)
  3. Firewall (nftables, nur in diesem Container):
       raus durch den Tunnel: alles; rein aus dem Tunnel: nur Antworten
       raus ins ARIA-Netz/Internet: nur zum RVS und zum Tunnel-Endpunkt
  4. satellite.py starten (SAT_KIND=vpn, Scan = SCAN_NETS aktiv abklopfen)
  5. Tunnel beobachten: status.json + events.jsonl (verbunden/getrennt, Dauer)
     fuer das ℹ-Fenster; Endpunkt mit Hostnamen (DynDNS) neu aufloesen

Env: VPN_ID, VPN_NAME, RVS_HOST/RVS_PORT/RVS_SNI (+ alles fuer satellite.py)
"""

from __future__ import annotations

import ipaddress
import json
import logging
import os
import re
import signal
import socket
import subprocess
import sys
import threading
import time

logging.basicConfig(level=logging.INFO, format="%(asctime)s [vpn] %(levelname)s %(message)s")
log = logging.getLogger("vpn")

VPN_DIR = os.environ.get("VPN_DIR", "/vpn")
CONF = os.path.join(VPN_DIR, "wg.conf")
STATUS = os.path.join(VPN_DIR, "status.json")
EVENTS = os.path.join(VPN_DIR, "events.jsonl")
IFACE = "wg0"
HANDSHAKE_OK_S = 180          # juenger = verbunden (WireGuard erneuert alle 2 min)
WATCH_S = 15

_stop = threading.Event()
_conn = {"connected": False, "since": 0, "rx": 0, "tx": 0}   # fuer "getrennt" beim Stoppen


def sh(*args: str, check: bool = True, inp: str | None = None) -> str:
    p = subprocess.run(list(args), capture_output=True, text=True, input=inp)
    if check and p.returncode != 0:
        raise RuntimeError(f"{' '.join(args[:4])}: {(p.stderr or p.stdout).strip()[:300]}")
    return p.stdout


def event(kind: str, **extra) -> None:
    rec = {"ts": int(time.time()), "event": kind, **extra}
    try:
        with open(EVENTS, "a", encoding="utf-8") as f:
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
    except OSError as e:
        log.warning("events.jsonl nicht schreibbar: %s", e)


# ── WireGuard-Datei lesen ───────────────────────────────────────

def parse_conf(text: str) -> dict:
    """[Interface] + beliebig viele [Peer]. Schluessel ohne Gross/Klein."""
    iface: dict = {}
    peers: list = []
    cur = None
    for raw in text.splitlines():
        line = raw.split("#", 1)[0].strip()
        if not line:
            continue
        sec = re.match(r"^\[(\w+)\]$", line)
        if sec:
            name = sec.group(1).lower()
            if name == "interface":
                cur = iface
            elif name == "peer":
                cur = {}
                peers.append(cur)
            else:
                cur = None
            continue
        if cur is None or "=" not in line:
            continue
        k, v = (s.strip() for s in line.split("=", 1))
        k = k.lower()
        if k in cur and k in ("address", "allowedips", "dns"):
            cur[k] += "," + v
        else:
            cur[k] = v
    if not iface.get("privatekey"):
        raise ValueError("[Interface] PrivateKey fehlt")
    if not peers or not all(p.get("publickey") for p in peers):
        raise ValueError("[Peer] PublicKey fehlt")
    return {"iface": iface, "peers": peers}


def _split(v: str) -> list[str]:
    return [x.strip() for x in (v or "").split(",") if x.strip()]


def setconf_text(c: dict) -> str:
    """Nur was `wg setconf` versteht (Address/DNS/MTU/PostUp sind wg-quick)."""
    out = ["[Interface]", f"PrivateKey = {c['iface']['privatekey']}"]
    if c["iface"].get("listenport"):
        out.append(f"ListenPort = {c['iface']['listenport']}")
    for p in c["peers"]:
        out += ["", "[Peer]", f"PublicKey = {p['publickey']}"]
        for k, name in (("presharedkey", "PresharedKey"), ("endpoint", "Endpoint"),
                        ("allowedips", "AllowedIPs")):
            if p.get(k):
                out.append(f"{name} = {p[k]}")
        out.append(f"PersistentKeepalive = {p.get('persistentkeepalive') or 25}")
    return "\n".join(out) + "\n"


def endpoint_host_port(ep: str) -> tuple[str, int]:
    m = re.match(r"^\[?([^\]]+?)\]?:(\d+)$", ep.strip())
    if not m:
        raise ValueError(f"Endpoint unlesbar: {ep}")
    return m.group(1), int(m.group(2))


def resolve4(host: str) -> list[str]:
    try:
        ipaddress.ip_address(host)
        return [host]
    except ValueError:
        pass
    try:
        return sorted({a[4][0] for a in socket.getaddrinfo(host, None, socket.AF_INET)})
    except OSError:
        return []


# ── Netz + Firewall ─────────────────────────────────────────────

def default_gw() -> str:
    m = re.search(r"default via (\S+)", sh("ip", "-4", "route", "show", "default"))
    if not m:
        raise RuntimeError("Container hat kein Default-Gateway")
    return m.group(1)


def firewall(rvs_ips: list[str], rvs_port: int, ep_ips: list[str], listen: int = 0) -> None:
    rvs = ", ".join(rvs_ips) if rvs_ips else "127.0.0.2"
    eps = ", ".join(ep_ips) if ep_ips else "127.0.0.2"
    # Server-Verbindung: der Verteiler (vpn_relay.py) reicht die Pakete der
    # Gegenseite an diesen Port weiter — nur WireGuard, nur ueber eth0.
    lin = f'iifname "eth0" udp dport {listen} accept' if listen else ""
    lout = f'oifname "eth0" udp sport {listen} accept' if listen else ""
    # Nur die eigene Tabelle — "flush ruleset" wuerde Dockers NAT-Regeln fuer
    # den eingebauten DNS (127.0.0.11) im Container mit loeschen.
    sh("nft", "delete", "table", "inet", "vpnsat", check=False)
    rules = f"""
table inet vpnsat {{
  set endpoints {{ type ipv4_addr; elements = {{ {eps} }} }}
  chain input {{
    type filter hook input priority 0; policy drop;
    iif lo accept
    ct state established,related accept
    iifname "eth0" ip saddr @endpoints meta l4proto udp accept
    {lin}
  }}
  chain forward {{
    type filter hook forward priority 0; policy drop;
  }}
  chain output {{
    type filter hook output priority 0; policy drop;
    oif lo accept
    ct state established,related accept
    oifname "{IFACE}" accept
    oifname "eth0" ip daddr {{ {rvs} }} tcp dport {rvs_port} accept
    oifname "eth0" ip daddr @endpoints meta l4proto udp accept
    {lout}
  }}
}}
"""
    sh("nft", "-f", "-", inp=rules)
    log.info("Firewall aktiv: Tunnel nur ausgehend; ins ARIA-Netz nur RVS %s:%d", rvs, rvs_port)


def setup(c: dict) -> dict:
    rvs_host = os.environ.get("RVS_HOST", "")
    rvs_port = int(os.environ.get("RVS_PORT", "443") or 443)
    rvs_ips = resolve4(rvs_host)
    if not rvs_ips:
        raise RuntimeError(f"RVS {rvs_host!r} nicht aufloesbar")
    gw = default_gw()

    ep_ips: list[str] = []
    for p in c["peers"]:
        if p.get("endpoint"):
            h, _ = endpoint_host_port(p["endpoint"])
            ep_ips += resolve4(h)

    sh("ip", "link", "del", IFACE, check=False)
    try:
        sh("ip", "link", "add", IFACE, "type", "wireguard")
    except RuntimeError as e:
        raise RuntimeError(f"WireGuard im Kernel nicht verfuegbar (modprobe wireguard auf dem Host?): {e}")
    sh("wg", "setconf", IFACE, "/dev/stdin", inp=setconf_text(c))
    for a in _split(c["iface"].get("address", "")):
        if ":" not in a:
            sh("ip", "-4", "address", "add", a, "dev", IFACE)
    sh("ip", "link", "set", IFACE, "mtu", c["iface"].get("mtu") or "1420", "up")

    # RVS und Endpunkte immer am Tunnel vorbei
    for ip in set(rvs_ips + ep_ips):
        sh("ip", "route", "replace", f"{ip}/32", "via", gw, "dev", "eth0")
    nets = []
    for p in c["peers"]:
        for n in _split(p.get("allowedips", "")):
            if ":" in n:
                continue
            net = ipaddress.ip_network(n, strict=False)
            if net.prefixlen == 0:
                for half in ("0.0.0.0/1", "128.0.0.0/1"):
                    sh("ip", "route", "replace", half, "dev", IFACE)
            else:
                sh("ip", "route", "replace", str(net), "dev", IFACE)
            nets.append(str(net))
    firewall(rvs_ips, rvs_port, ep_ips, int(c["iface"].get("listenport") or 0))
    log.info("Tunnel %s oben: Netze %s, Endpunkt(e) %s", IFACE, ", ".join(nets) or "-", ", ".join(ep_ips) or "-")
    return {"nets": nets, "endpoints": ep_ips}


def scan_nets(c: dict) -> str:
    """Was der Satellit abklopft: SCAN_NETS (aus der Diagnostic) oder die
    AllowedIPs ohne 0.0.0.0/0. Zu grosse Netze kuerzt satellite.py selbst."""
    if os.environ.get("SCAN_NETS", "").strip():
        return os.environ["SCAN_NETS"].strip()
    nets = []
    for p in c["peers"]:
        for n in _split(p.get("allowedips", "")):
            if ":" not in n and not n.endswith("/0"):
                nets.append(n)
    return ",".join(nets)


# ── Beobachten ──────────────────────────────────────────────────

def watch(c: dict) -> None:
    connected = False
    since = 0
    last_resolve = 0.0
    while not _stop.is_set():
        try:
            dump = sh("wg", "show", IFACE, "dump").splitlines()[1:]
            now = int(time.time())
            hs = max((int(l.split("\t")[4]) for l in dump if l.count("\t") >= 6), default=0)
            rx = sum(int(l.split("\t")[5]) for l in dump if l.count("\t") >= 6)
            tx = sum(int(l.split("\t")[6]) for l in dump if l.count("\t") >= 6)
            ep = ", ".join(l.split("\t")[2] for l in dump if l.count("\t") >= 6)
            ok = hs > 0 and now - hs < HANDSHAKE_OK_S
            if ok and not connected:
                since = now
                event("verbunden", endpoint=ep)
                log.info("verbunden (%s)", ep)
            elif connected and not ok:
                event("getrennt", dauer_s=now - since, rx=rx, tx=tx)
                log.info("getrennt nach %d s", now - since)
            connected = ok
            _conn.update(connected=ok, since=since, rx=rx, tx=tx)
            st = {"state": "verbunden" if ok else "getrennt", "since": since if ok else 0,
                  "handshake": hs, "rx": rx, "tx": tx, "endpoint": ep, "updated": now}
            tmp = STATUS + ".tmp"
            with open(tmp, "w") as f:
                json.dump(st, f)
            os.replace(tmp, STATUS)
            # DynDNS beim Kunden: ohne Verbindung den Namen neu aufloesen
            if not ok and time.time() - last_resolve > 120:
                last_resolve = time.time()
                for p in c["peers"]:
                    if not p.get("endpoint"):
                        continue
                    h, port = endpoint_host_port(p["endpoint"])
                    for ip in resolve4(h):
                        sh("nft", "add", "element", "inet", "vpnsat", "endpoints", f"{{ {ip} }}", check=False)
                        sh("ip", "route", "replace", f"{ip}/32", "via", default_gw(), "dev", "eth0", check=False)
                    sh("wg", "set", IFACE, "peer", p["publickey"], "endpoint", f"{h}:{port}", check=False)
        except Exception as e:
            log.warning("Beobachten: %s", e)
        _stop.wait(WATCH_S)


def main() -> None:
    def _term(*_):
        _stop.set()
    signal.signal(signal.SIGTERM, _term)
    signal.signal(signal.SIGINT, _term)

    try:
        c = parse_conf(open(CONF, encoding="utf-8").read())
        info = setup(c)
    except Exception as e:
        log.error("Tunnel nicht aufgebaut: %s", e)
        event("fehler", meldung=str(e)[:300])
        with open(STATUS, "w") as f:
            json.dump({"state": "fehler", "error": str(e)[:300], "updated": int(time.time())}, f)
        time.sleep(60)          # nicht im Sekundentakt neu starten
        sys.exit(1)
    event("gestartet", netze=info["nets"])

    env = dict(os.environ)
    env.update({
        "SAT_KIND": "vpn",
        "SATELLITE_ID": env.get("SATELLITE_ID") or f"vpn-{env.get('VPN_ID', 'x')}",
        "SATELLITE_LOCATION": env.get("VPN_NAME") or env.get("SATELLITE_LOCATION") or "VPN",
        "SCAN_NETS": scan_nets(c),
        "LIST_SELF": "false",
        "CONTROL_ENABLED": env.get("CONTROL_ENABLED", "true"),
    })
    threading.Thread(target=watch, args=(c,), daemon=True).start()
    proc = None
    while not _stop.is_set():
        proc = subprocess.Popen([sys.executable, "-u", "/app/satellite.py"], env=env)
        while proc.poll() is None and not _stop.is_set():
            time.sleep(1)
        if not _stop.is_set():
            log.warning("satellite.py beendet (rc=%s) — starte neu", proc.returncode)
            _stop.wait(5)
    if proc and proc.poll() is None:
        proc.terminate()
        try:
            proc.wait(10)
        except subprocess.TimeoutExpired:
            proc.kill()
    if _conn["connected"]:
        event("getrennt", dauer_s=int(time.time()) - _conn["since"], rx=_conn["rx"], tx=_conn["tx"])
    event("gestoppt")
    sh("ip", "link", "del", IFACE, check=False)


if __name__ == "__main__":
    main()