#!/usr/bin/env python3
"""mock_wg_endpoint.py — мок WireGuard/AmneziaWG-endpoint для netns-стенда (задача #21).

Зачем: интеграционные тесты split-tunnel должны идти без реального VPN-провайдера и
без настоящего туннеля. Мок слушает UDP-порт 51820 в netns `uplink` («ISP»),
распознаёт тип сообщения WireGuard по первому байту и отвечает на handshake
initiation (type 1, 148 байт) корректным по структуре handshake response
(type 2, 92 байта), подставляя sender_index клиента как receiver_index.

ЧЕГО МОК НЕ ДЕЛАЕТ (важно, чтобы не принять его за настоящий peer):
  * не выполняет Noise IK handshake и никакой криптографии — поля «зашифрованы»
    случайными байтами;
  * не маршрутизирует трафик и не поднимает TUN-интерфейс.
Поэтому мок НЕ подходит для проверки шифрования/MTU/throughput — только для
проверки топологии, policy routing, метки трафика и fail-closed сценариев.

Режим работы переключается на лету через файл `--mode-file`:
  respond — отвечать на handshake initiation (провайдер доступен);
  drop    — молча считать пакеты и не отвечать (имитация «провайдер недоступен»).

Лог: JSONL (по строке на событие), пишется и сбрасывается на диск сразу, чтобы
тесты могли его читать без гонок. Файл `--ready` создаётся после bind().

Только stdlib — запускается и на хосте, и в alpine-образе приложения.
"""

from __future__ import annotations

import argparse
import json
import os
import random
import socket
import struct
import sys
import time

MSG_HANDSHAKE_INITIATION = 1
MSG_HANDSHAKE_RESPONSE = 2
MSG_COOKIE_REPLY = 3
MSG_TRANSPORT_DATA = 4

TYPE_NAMES = {
    MSG_HANDSHAKE_INITIATION: "handshake_initiation",
    MSG_HANDSHAKE_RESPONSE: "handshake_response",
    MSG_COOKIE_REPLY: "cookie_reply",
    MSG_TRANSPORT_DATA: "transport_data",
}

HANDSHAKE_INITIATION_LEN = 148
HANDSHAKE_RESPONSE_LEN = 92


def msg_type(payload: bytes) -> int:
    """WireGuard message type: младший байт little-endian uint32 (3 байта reserved)."""
    if len(payload) < 4:
        return -1
    return payload[0]


def build_handshake_response(initiation: bytes) -> bytes:
    """Собрать структурно корректный type 2 (92 байта) в ответ на type 1.

    Раскладка: type(1)+reserved(3) sender(4) receiver(4) ephemeral(32)
               encrypted_nothing(16) mac1(16) mac2(16) = 92
    """
    sender_index = initiation[4:8] if len(initiation) >= 8 else b"\x00" * 4
    out = bytearray(HANDSHAKE_RESPONSE_LEN)
    out[0] = MSG_HANDSHAKE_RESPONSE
    # sender_index мока — случайный, receiver_index = sender_index клиента
    out[4:8] = struct.pack("<I", random.getrandbits(32))
    out[8:12] = sender_index
    out[12:44] = os.urandom(32)   # unencrypted_ephemeral
    out[44:60] = os.urandom(16)   # encrypted_nothing + poly1305 tag
    out[60:76] = os.urandom(16)   # mac1
    out[76:92] = os.urandom(16)   # mac2
    return bytes(out)


class JsonlLog:
    """Пишущий JSONL-лог с немедленным flush (тесты читают файл параллельно)."""

    def __init__(self, path: str | None) -> None:
        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))
        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 read_mode(mode_file: str | None, default: str) -> str:
    """Режим читается из файла на каждом пакете: тесты переключают без рестарта."""
    if not mode_file:
        return default
    try:
        with open(mode_file, encoding="utf-8") as fh:
            mode = fh.read().strip()
    except OSError:
        return default
    return mode if mode in ("respond", "drop") else default


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description="Мок WireGuard-endpoint для netns-стенда (#21)")
    ap.add_argument("--bind", default="0.0.0.0:51820", help="адрес:порт для прослушивания UDP")
    ap.add_argument("--log", default=os.environ.get("AWG_MOCK_LOG", ""), help="JSONL-лог событий")
    ap.add_argument("--ready", default=os.environ.get("AWG_MOCK_READY", ""), help="файл-флаг готовности")
    ap.add_argument(
        "--mode-file",
        default=os.environ.get("AWG_MOCK_MODE_FILE", ""),
        help="файл с режимом respond|drop (переключение на лету)",
    )
    ap.add_argument("--mode", default=os.environ.get("AWG_MOCK_MODE", "respond"), choices=["respond", "drop"])
    args = ap.parse_args(argv)

    host, _, port_s = args.bind.rpartition(":")
    host = host or "0.0.0.0"
    port = int(port_s or "51820")

    log = JsonlLog(args.log or None)
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    try:
        sock.bind((host, port))
    except OSError as exc:
        log.write(event="bind_error", addr=f"{host}:{port}", error=str(exc))
        print(f"mock-endpoint: bind {host}:{port} failed: {exc}", file=sys.stderr)
        return 4
    sock.settimeout(0.25)

    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"{host}:{port}", pid=os.getpid(), mode=args.mode)
    print(f"mock-endpoint: udp/{port} listen, mode={args.mode}", flush=True)

    rx_total = 0
    tx_total = 0
    try:
        while True:
            try:
                payload, peer = sock.recvfrom(65535)
            except socket.timeout:
                continue
            except OSError:
                break

            rx_total += 1
            mtype = msg_type(payload)
            mode = read_mode(args.mode_file or None, args.mode)
            log.write(
                event="rx",
                msg_type=mtype,
                msg_name=TYPE_NAMES.get(mtype, "unknown"),
                length=len(payload),
                src=f"{peer[0]}:{peer[1]}",
                rx_total=rx_total,
                mode=mode,
            )

            if mode == "drop":
                continue
            if mtype != MSG_HANDSHAKE_INITIATION:
                log.write(event="ignored", reason="not_handshake_initiation", msg_type=mtype)
                continue
            if len(payload) != HANDSHAKE_INITIATION_LEN:
                log.write(
                    event="ignored",
                    reason="bad_initiation_length",
                    length=len(payload),
                    expected=HANDSHAKE_INITIATION_LEN,
                )
                continue

            reply = build_handshake_response(payload)
            sock.sendto(reply, peer)
            tx_total += 1
            log.write(
                event="tx",
                msg_type=MSG_HANDSHAKE_RESPONSE,
                length=len(reply),
                dst=f"{peer[0]}:{peer[1]}",
                tx_total=tx_total,
            )
    except KeyboardInterrupt:
        pass
    finally:
        sock.close()
        log.write(event="stop", rx_total=rx_total, tx_total=tx_total)
        log.close()
        if args.ready and os.path.exists(args.ready):
            os.unlink(args.ready)
    return 0


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