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

Запускается внутри конкретного netns (`ip netns exec <ns> python3 probe.py ...`)
и умеет:

  * `--wg-initiation` (по умолчанию) — отправить WireGuard handshake initiation
    (type 1, 148 байт, раскладка как в протоколе) и дождаться ответа мока;
  * `--raw TEXT` — отправить произвольный UDP-datagram (для проверок утечки /
    kill-switch), по умолчанию ответа не ожидает.

Результат печатается одной JSON-строкой в stdout, код возврата:
  0 — ожидаемое поведение достигнуто, 1 — нет, 2 — ошибка аргументов/сокета.
"""

from __future__ import annotations

import argparse
import json
import os
import struct
import sys
import time

MSG_HANDSHAKE_INITIATION = 1
MSG_HANDSHAKE_RESPONSE = 2
HANDSHAKE_INITIATION_LEN = 148


def build_handshake_initiation(sender_index: int | None = None) -> bytes:
    """148 байт: type(1)+reserved(3) sender(4) ephemeral(32) enc_static(48) enc_ts(28) mac1(16) mac2(16)."""
    if sender_index is None:
        sender_index = struct.unpack("<I", os.urandom(4))[0]
    out = bytearray(HANDSHAKE_INITIATION_LEN)
    out[0] = MSG_HANDSHAKE_INITIATION
    out[4:8] = struct.pack("<I", sender_index)
    out[8:40] = os.urandom(32)    # unencrypted_ephemeral
    out[40:88] = os.urandom(48)   # encrypted_static (32) + poly1305 tag (16)
    out[88:116] = os.urandom(28)  # encrypted_timestamp (12) + tag (16)
    out[116:132] = os.urandom(16)  # mac1
    out[132:148] = os.urandom(16)  # mac2
    return bytes(out)


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description="Тестовый трафик для netns-стенда (#21)")
    ap.add_argument("--dst", required=True, help="IP назначения")
    ap.add_argument("--port", type=int, required=True, help="UDP-порт назначения")
    ap.add_argument("--timeout", type=float, default=2.0, help="ожидание ответа, с")
    ap.add_argument("--raw", default=None, help="отправить произвольный текст вместо handshake initiation")
    ap.add_argument("--expect-timeout", action="store_true", help="ожидать отсутствия ответа (для kill-switch/drop)")
    ap.add_argument("--source-port", type=int, default=0, help="локальный порт (0 — случайный)")
    args = ap.parse_args(argv)

    if args.raw is None:
        payload = build_handshake_initiation()
        expect_reply = not args.expect_timeout
        kind = "wg_handshake_initiation"
    else:
        payload = args.raw.encode("utf-8")
        # для raw-трафика ответ не ожидается, если явно не сказано обратное
        expect_reply = not args.expect_timeout
        kind = "raw_udp"

    import socket

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    if args.source_port:
        sock.bind(("", args.source_port))
    sock.settimeout(args.timeout)

    started = time.time()
    try:
        sent = sock.sendto(payload, (args.dst, args.port))
    except OSError as exc:
        print(json.dumps({"ok": False, "kind": kind, "error": f"sendto: {exc}"}))
        return 2

    reply = None
    try:
        data, peer = sock.recvfrom(65535)
        reply = {
            "type": data[0] if data else -1,
            "length": len(data or b""),
            "src": f"{peer[0]}:{peer[1]}",
            "receiver_index": (
                struct.unpack("<I", data[8:12])[0] if len(data) >= 12 and data[0] == MSG_HANDSHAKE_RESPONSE else None
            ),
            "sender_index_echo": struct.unpack("<I", payload[4:8])[0],
        }
    except socket.timeout:
        reply = None
    finally:
        sock.close()

    got_reply = reply is not None
    ok = got_reply if expect_reply else (not got_reply)
    out = {
        "ok": ok,
        "kind": kind,
        "sent": sent,
        "dst": f"{args.dst}:{args.port}",
        "expect_reply": expect_reply,
        "reply_received": got_reply,
        "reply": reply,
        "rtt_ms": round((time.time() - started) * 1000, 2),
    }
    print(json.dumps(out, sort_keys=True))
    return 0 if ok else 1


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