#!/usr/bin/env python3
"""leak_sink.py — детектор утечки трафика для netns-стенда (задача #21).

Слушает UDP и TCP порт 9999 на «публичном» адресе 198.51.100.9 внутри netns
`uplink` (адрес из TEST-NET-2, RFC 5737). Если при включённом kill-switch
шлюза сюда пришёл хоть один пакет — значит трафик, предназначенный для
туннеля, утёк напрямую: тест падает.

Лог — JSONL, пишется немедленно (тесты читают файл, не дожидаясь остановки).
"""

from __future__ import annotations

import argparse
import json
import os
import socket
import sys
import threading
import time


class JsonlLog:
    def __init__(self, path: str | None) -> None:
        self._lock = threading.Lock()
        self._fh = None
        if path:
            os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
            self._fh = open(path, "a", buffering=1, encoding="utf-8")

    def write(self, **event: object) -> None:
        if self._fh is None:
            return
        event.setdefault("ts", round(time.time(), 6))
        with self._lock:
            self._fh.write(json.dumps(event, sort_keys=True) + "\n")
            self._fh.flush()

    def close(self) -> None:
        if self._fh is not None:
            self._fh.close()
            self._fh = None


def udp_loop(bind: tuple[str, int], log: JsonlLog, stop: threading.Event) -> None:
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(bind)
    sock.settimeout(0.25)
    while not stop.is_set():
        try:
            payload, peer = sock.recvfrom(65535)
        except socket.timeout:
            continue
        except OSError:
            break
        log.write(event="leak", proto="udp", length=len(payload), src=f"{peer[0]}:{peer[1]}")
    sock.close()


def tcp_loop(bind: tuple[str, int], log: JsonlLog, stop: threading.Event) -> None:
    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(bind)
    srv.listen(8)
    srv.settimeout(0.25)
    while not stop.is_set():
        try:
            conn, peer = srv.accept()
        except socket.timeout:
            continue
        except OSError:
            break
        log.write(event="leak", proto="tcp", src=f"{peer[0]}:{peer[1]}")
        conn.close()
    srv.close()


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description="Детектор утечки трафика для netns-стенда (#21)")
    ap.add_argument("--ip", default=os.environ.get("AWG_LEAK_IP", "198.51.100.9"))
    ap.add_argument("--port", type=int, default=int(os.environ.get("AWG_LEAK_PORT", "9999")))
    ap.add_argument("--log", default=os.environ.get("AWG_LEAK_LOG", ""))
    ap.add_argument("--ready", default=os.environ.get("AWG_LEAK_READY", ""))
    args = ap.parse_args(argv)

    bind = (args.ip, args.port)
    log = JsonlLog(args.log or None)
    stop = threading.Event()

    udp = threading.Thread(target=udp_loop, args=(bind, log, stop), daemon=True)
    tcp = threading.Thread(target=tcp_loop, args=(bind, log, stop), daemon=True)
    udp.start()
    tcp.start()

    if args.ready:
        os.makedirs(os.path.dirname(os.path.abspath(args.ready)) or ".", exist_ok=True)
        with open(args.ready, "w", encoding="utf-8") as fh:
            fh.write(f"{os.getpid()}\n")

    log.write(event="listen", addr=f"{args.ip}:{args.port}", pid=os.getpid())
    print(f"leak-sink: udp+tcp {args.ip}:{args.port} listen", flush=True)
    try:
        while True:
            time.sleep(0.5)
    except KeyboardInterrupt:
        pass
    finally:
        stop.set()
        udp.join(timeout=2)
        tcp.join(timeout=2)
        log.write(event="stop")
        log.close()
        if args.ready and os.path.exists(args.ready):
            os.unlink(args.ready)
    return 0


if __name__ == "__main__":
    sys.exit(main())