"""
ARIA VPN-Verteiler — EIN UDP-Port fuer alle WireGuard-Server-Verbindungen.

Kunden waehlen sich alle auf denselben Port ein (UDP 443 am RVS-Hostnamen,
per NPM-Stream bzw. Router hierher weitergeleitet). Jede Server-Verbindung
laeuft aber in ihrem EIGENEN Container (eigener Netz-Namensraum, wegen
ueberlappender Kundennetze). Der Verteiler ordnet die Pakete zu, ohne etwas
zu entschluesseln:

  Typ 1 (Handshake-Anfang): mac1 = BLAKE2s(Schluessel aus dem oeffentlichen
        Schluessel des EMPFAENGERS, Nachricht). Jede Verbindung hat ein eigenes
        Server-Schluesselpaar → der Verteiler rechnet mac1 fuer jede bekannte
        Verbindung nach und weiss, wohin das Paket gehoert.
  Typ 2 (Antwort des Containers): merkt sich dessen Sitzungsnummer.
  Typ 3/4 (Cookie/Daten): nach Absender-Adresse, sonst nach Sitzungsnummer
        (Kunde hat die Adresse gewechselt).

Jeder Absender bekommt einen eigenen Upstream-Socket zum Container, so finden
die Antworten zurueck. Unbekanntes wird still verworfen — wie WireGuard selbst
(ein Portscan sieht nichts).

Laeuft mit network_mode: host (erreicht die Container-IPs im Netz aria-vpn,
dort ist Container-zu-Container gesperrt). Liest /vpn/relay.json, das die
Diagnostic schreibt: [{id, ip, port, pub}].
"""

from __future__ import annotations

import asyncio
import base64
import hashlib
import json
import logging
import os
import struct
import time

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

LISTEN_PORT = int(os.environ.get("VPN_RELAY_PORT", "51820") or 51820)
MAP_FILE = os.environ.get("VPN_RELAY_MAP", "/vpn/relay.json")
IDLE_S = 600
MAX_CLIENTS = 2048

_targets: list[dict] = []        # [{id, addr:(ip,port), mac1key}]
_map_mtime = 0.0
_by_index: dict[int, str] = {}   # Sitzungsnummer des Containers -> id
_clients: dict = {}              # Kunden-Adresse -> Upstream


def _mac1_key(pub_b64: str) -> bytes:
    pub = base64.b64decode(pub_b64)
    return hashlib.blake2s(b"mac1----" + pub).digest()


def load_map() -> None:
    global _targets, _map_mtime
    try:
        m = os.path.getmtime(MAP_FILE)
    except OSError:
        if _targets:
            log.info("keine Zuordnung mehr (%s fehlt)", MAP_FILE)
        _targets, _map_mtime = [], 0.0
        return
    if m == _map_mtime:
        return
    try:
        rows = json.load(open(MAP_FILE))
        _targets = [{"id": r["id"], "addr": (r["ip"], int(r.get("port", 51820))),
                     "mac1key": _mac1_key(r["pub"])} for r in rows if r.get("ip") and r.get("pub")]
        _map_mtime = m
        log.info("%d Server-Verbindung(en): %s", len(_targets),
                 ", ".join(f"{t['id']}→{t['addr'][0]}" for t in _targets) or "-")
    except Exception as e:
        log.warning("relay.json unlesbar: %s", e)


def target_for_init(data: bytes) -> dict | None:
    if len(data) != 148:
        return None
    for t in _targets:
        mac = hashlib.blake2s(data[:116], key=t["mac1key"], digest_size=16).digest()
        if mac == data[116:132]:
            return t
    return None


class Upstream(asyncio.DatagramProtocol):
    """Ein Socket je Kunde → Container; leitet Antworten zurueck."""

    def __init__(self, front, client_addr, target):
        self.front, self.client, self.target = front, client_addr, target
        self.transport = None
        self.last = time.time()

    def connection_made(self, transport):
        self.transport = transport

    def datagram_received(self, data, addr):
        self.last = time.time()
        if len(data) >= 8 and data[0] == 2:          # Antwort: Sitzungsnummer merken
            _by_index[struct.unpack_from("<I", data, 4)[0]] = self.target["id"]
        self.front.sendto(data, self.client)

    def error_received(self, exc):
        pass


class Front(asyncio.DatagramProtocol):
    def connection_made(self, transport):
        self.transport = transport

    def datagram_received(self, data, addr):
        asyncio.ensure_future(self.handle(data, addr))

    async def handle(self, data: bytes, addr) -> None:
        if len(data) < 4:
            return
        load_map()
        kind = data[0]
        up = _clients.get(addr)
        target = None
        if kind == 1:
            target = target_for_init(data)
            if not target:
                return                                  # fremder Schluessel → still verwerfen
        elif up is None and kind == 4 and len(data) >= 8:
            tid = _by_index.get(struct.unpack_from("<I", data, 4)[0])
            target = next((t for t in _targets if t["id"] == tid), None)
            if not target:
                return
        elif up is None:
            return
        if up and target and up.target["id"] != target["id"]:
            up.transport.close()
            up = None
        if up is None or up.target["addr"] != (target or up.target)["addr"]:
            if up:
                up.transport.close()
            if len(_clients) >= MAX_CLIENTS:
                return
            loop = asyncio.get_running_loop()
            t = target or up.target
            _, up = await loop.create_datagram_endpoint(
                lambda: Upstream(self.transport, addr, t), remote_addr=t["addr"])
            _clients[addr] = up
        up.last = time.time()
        up.transport.sendto(data)


async def janitor() -> None:
    while True:
        await asyncio.sleep(60)
        now = time.time()
        for a, up in list(_clients.items()):
            if now - up.last > IDLE_S:
                up.transport.close()
                del _clients[a]
        if len(_by_index) > 8192:
            _by_index.clear()


async def main() -> None:
    load_map()
    loop = asyncio.get_running_loop()
    await loop.create_datagram_endpoint(Front, local_addr=("0.0.0.0", LISTEN_PORT))
    log.info("lausche auf UDP %d", LISTEN_PORT)
    await janitor()


if __name__ == "__main__":
    asyncio.run(main())