#!/usr/bin/env python3
"""Drive synthetic TCP/UDP traffic against a game's actual /32 IPs from
the live GPNC catalog, then verify the detector fired.

Pulls game IPs + port rules from GPNC's /api/catalog so the test stays in
sync with whatever's currently loaded. Runs from the SUBSCRIBER side of
the bridge — the source MAC must be on the subscriber-side L2 segment for
the flow detector to consider it a legitimate trigger (infra MACs are
skipped).

Usage:
    # Default — drive WoT for 30s with TCP+UDP at 10pps each
    python traffic_gen.py --game "World of Tanks"

    # UDP only, faster, against a specific IP from the game's set
    python traffic_gen.py --game "Albion Online" --proto udp --rate 50 \
        --duration 60 --ip 5.188.125.14

    # List what games are detectable
    python traffic_gen.py --list

The script:
  1. Pulls /api/catalog from gpnc-status, finds the named game
  2. Reads its /32 IPs (first one or --ip override) + a port from its rules
  3. Sends TCP SYNs and/or UDP datagrams at --rate pps for --duration s
  4. Polls /api/status before+after to see if a session launched
  5. Pulls the journal and reports any "L2-inline launch" lines for our MAC

Notes:
  - Non-root mode uses Python sockets — TCP SYNs become full connect()
    attempts, UDP is just sendto(). For SYN floods at >100pps, run with
    sudo to use raw sockets (requires root + scapy if --raw is passed).
  - If the IP is unreachable (no route), TCP connect() fails fast, but
    the SYN still leaves the box and that's all we need for the detector.
"""
from __future__ import annotations

import argparse
import json
import socket
import subprocess
import sys
import time
import urllib.request
from typing import Any

DEFAULT_STATUS = "http://10.200.90.14"  # gpnc-status reverse-proxy in NPM
DEFAULT_GPNC_API = "http://10.80.4.15:50055"  # gpnc HTTP (direct)


def fetch_catalog(base: str) -> dict[str, Any]:
    with urllib.request.urlopen(f"{base}/api/catalog", timeout=10) as r:
        return json.loads(r.read().decode())


def fetch_status(base: str) -> dict[str, Any]:
    with urllib.request.urlopen(f"{base}/api/status", timeout=10) as r:
        return json.loads(r.read().decode())


def fetch_game(base: str, config_id: str) -> dict[str, Any]:
    with urllib.request.urlopen(
        f"{base}/api/catalog/{config_id}", timeout=10) as r:
        return json.loads(r.read().decode())


def list_games(base: str) -> None:
    cat = fetch_catalog(base)
    rows = [g for g in cat.get("games", [])
            if g.get("signature") in ("ip32", "ip32+port")]
    rows.sort(key=lambda g: -(g.get("ip32_count") or 0))
    print(f"{len(rows)} ip32-detectable games:")
    print(f"  {'Game':30s} {'/32s':>4s} {'ports':>5s}  IPs")
    for g in rows:
        ips = g.get("ip32s") or []
        sample = ", ".join(ips[:3]) + (f" (+{len(ips)-3})" if len(ips) > 3 else "")
        print(f"  {g['name'][:30]:30s} {g.get('ip32_count', 0):>4d} "
              f"{g.get('port_rule_count', 0):>5d}  {sample}")


def find_game(catalog: dict, name: str) -> dict | None:
    name_lower = name.lower()
    # Exact match first, then prefix.
    for g in catalog.get("games", []):
        if g["name"].lower() == name_lower:
            return g
    for g in catalog.get("games", []):
        if g["name"].lower().startswith(name_lower):
            return g
    return None


def pick_port(entry: dict) -> tuple[int, str]:
    """Pick a representative port from the game's first non-skip rule.
    Returns (port, proto) where proto is 'tcp' or 'udp'.
    Falls back to (25000, 'udp') if no usable rule found."""
    cfg = entry.get("Config") or {}
    skip = {22, 25, 53, 80, 110, 123, 137, 138, 139, 143, 161, 162, 389,
            443, 445, 465, 514, 587, 636, 853, 993, 995, 1900, 3389,
            5060, 5061, 5222, 5228, 5353, 5355, 5938, 8080, 8443, 843}
    for it in cfg.get("Intercept") or []:
        if not isinstance(it, dict):
            continue
        if (it.get("Action") or "").lower() != "proxy":
            continue
        proto = (it.get("Protocol") or "udp").lower()
        if proto not in ("tcp", "udp"):
            proto = "udp"
        ports = (it.get("DestinationPort") or "").split(",")
        for p in ports:
            p = p.strip()
            if not p:
                continue
            # Take the first port in any range like "5000-6000".
            first = p.split("-")[0].split(":")[0].strip()
            try:
                pn = int(first)
            except ValueError:
                continue
            if pn in skip:
                continue
            if 1 <= pn <= 65535:
                return pn, proto
    return 25000, "udp"


def warm_arp(target_ip: str) -> None:
    """Force ARP resolution + first-hop reachability before the burst.
    On a cold subscriber (e.g. fresh netns), the first packet of any flow
    triggers a synchronous ARP request to the gateway and gets dropped if
    the response doesn't return inside the kernel's queue threshold. By
    sending a few priming packets and waiting briefly, we ensure the
    detector sees a stable stream instead of just a half-resolved ARP."""
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.settimeout(0.5)
        # 3 priming UDP packets to the target — the first ones trigger ARP,
        # by the third the cache is hot.
        for _ in range(3):
            try:
                s.sendto(b"\x00", (target_ip, 1))  # arbitrary port; we just want the ARP
            except OSError:
                pass
            time.sleep(0.3)
        s.close()
    except Exception:
        pass


def send_tcp(ip: str, port: int, count: int, gap: float) -> int:
    """Open and immediately close TCP connections to (ip, port). Each
    attempt sends at least one SYN — that's all the detector needs."""
    sent = 0
    for _ in range(count):
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            s.settimeout(0.2)
            try:
                s.connect((ip, port))
            except (socket.timeout, OSError):
                pass
            finally:
                s.close()
            sent += 1
        except Exception as e:
            print(f"  tcp send error: {e}", file=sys.stderr)
        if gap > 0:
            time.sleep(gap)
    return sent


def send_udp(ip: str, port: int, count: int, gap: float) -> int:
    """Send UDP datagrams to (ip, port). Payload is a tiny game-shaped
    packet (~32 bytes) to avoid looking like throughput traffic."""
    sent = 0
    payload = bytes(range(32))  # 32-byte deterministic payload
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        for _ in range(count):
            try:
                s.sendto(payload, (ip, port))
                sent += 1
            except OSError as e:
                # ENETUNREACH etc. — packet still left the host? No: sendto
                # fails before the packet hits the wire if the route is
                # unknown. Print and continue.
                print(f"  udp send error: {e}", file=sys.stderr)
            if gap > 0:
                time.sleep(gap)
    finally:
        s.close()
    return sent


def gpnc_journal_excerpt(host: str, since_iso: str, mac_hint: str) -> str:
    """Pull the last few seconds of the gpnc journal via SSH and grep for
    L2-inline launch lines for our MAC."""
    try:
        out = subprocess.check_output([
            "ssh",
            "-o", "BatchMode=yes",
            "-o", "ConnectTimeout=5",
            f"wtfast@{host}",
            f"sudo -n journalctl -u gpnc --since '{since_iso}' --no-pager "
            f"| grep -E 'L2-inline launch|launch ok' | tail -10",
        ], stderr=subprocess.DEVNULL, timeout=8.0).decode()
        return out.strip()
    except Exception:
        return "(could not fetch journal — pass --gpnc-host or check ssh key)"


def main() -> int:
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--status-url", default=DEFAULT_STATUS,
                   help=f"gpnc-status base URL (default {DEFAULT_STATUS})")
    p.add_argument("--gpnc-host", default="10.80.4.15",
                   help="gpnc VM IP for journal grep (default 10.80.4.15)")
    p.add_argument("--list", action="store_true",
                   help="list games with ip32 signature and exit")
    p.add_argument("--game", help="game name (case-insensitive, prefix ok)")
    p.add_argument("--ip", help="override IP (must be one from the game's /32 set)")
    p.add_argument("--port", type=int, help="override destination port")
    p.add_argument("--proto", choices=("tcp", "udp", "both"), default="both",
                   help="protocol(s) to drive (default both)")
    p.add_argument("--rate", type=int, default=20,
                   help="packets per second (default 20)")
    p.add_argument("--duration", type=int, default=90,
                   help="seconds to run (default 90 — long enough for ARP "
                        "warm + detector trigger + sustained flow)")
    p.add_argument("--no-verify", action="store_true",
                   help="skip post-run journal/status check")
    args = p.parse_args()

    if args.list:
        list_games(args.status_url)
        return 0

    if not args.game:
        print("error: --game required (use --list to see options)", file=sys.stderr)
        return 1

    cat = fetch_catalog(args.status_url)
    g = find_game(cat, args.game)
    if not g:
        print(f"error: no game matching {args.game!r}", file=sys.stderr)
        return 2
    if not g.get("ip32_count"):
        print(f"error: game {g['name']!r} has no /32 IPs — port-only "
              "games can't be tested with this script", file=sys.stderr)
        return 3

    detail = fetch_game(args.status_url, g["config_id"])
    ips = g.get("ip32s") or []
    if args.ip:
        if args.ip not in ips:
            print(f"warn: --ip {args.ip} is NOT in game's /32 set "
                  f"({len(ips)} known); proceeding anyway", file=sys.stderr)
        target_ip = args.ip
    else:
        target_ip = ips[0]

    if args.port:
        port = args.port
        proto_default = "udp"
    else:
        port, proto_default = pick_port(detail)

    print(f"=== traffic_gen.py ===")
    print(f"  game:     {g['name']} (configID={g['config_id']})")
    print(f"  ip32s:    {len(ips)} known, hitting {target_ip}")
    print(f"  port:     {port} ({proto_default} default)")
    print(f"  rate:     {args.rate} pps  duration: {args.duration}s")
    print(f"  proto:    {args.proto}")

    # Snapshot status before
    if not args.no_verify:
        before = fetch_status(args.status_url)
        before_pairs = before.get("active_pairs") or []
        print(f"  before:   {len(before_pairs)} active pairs")

    since = time.strftime("%Y-%m-%d %H:%M:%S")

    # Warm ARP cache + first-hop reachability before the main burst.
    # Cold subscribers (fresh netns, just-booted laptops) need this — the
    # very first packet of a flow gets dropped while the kernel's ARP
    # request is in flight to the gateway. Without warming, ~30% of test
    # runs against unfamiliar destinations result in "no detector trigger"
    # purely because the burst ended before any packet reached br-gpnc.
    print(f"  warming:  ARP + route to {target_ip}")
    warm_arp(target_ip)
    time.sleep(1.0)

    total_pkts = args.rate * args.duration
    gap = 1.0 / args.rate if args.rate > 0 else 0.0

    do_tcp = args.proto in ("tcp", "both") or (args.proto == "both" and proto_default == "tcp")
    do_udp = args.proto in ("udp", "both") or (args.proto == "both" and proto_default == "udp")
    if args.proto == "tcp":
        do_tcp, do_udp = True, False
    elif args.proto == "udp":
        do_tcp, do_udp = False, True

    sent_tcp = sent_udp = 0
    t0 = time.time()
    if do_tcp and do_udp:
        # Interleave — half the rate on each.
        per = total_pkts // 2
        sent_tcp = send_tcp(target_ip, port, per, gap * 2)
        sent_udp = send_udp(target_ip, port, per, gap * 2)
    elif do_tcp:
        sent_tcp = send_tcp(target_ip, port, total_pkts, gap)
    elif do_udp:
        sent_udp = send_udp(target_ip, port, total_pkts, gap)
    elapsed = time.time() - t0

    print(f"\n=== traffic complete ===")
    print(f"  elapsed:  {elapsed:.1f}s")
    print(f"  sent:     tcp={sent_tcp}  udp={sent_udp}  total={sent_tcp+sent_udp}")

    if args.no_verify:
        return 0

    # Give the controller a beat to register the launch.
    time.sleep(2)
    after = fetch_status(args.status_url)
    after_pairs = after.get("active_pairs") or []
    new_pairs = [p for p in after_pairs
                 if (p.get("config_id") == g["config_id"])
                 and p not in before_pairs]
    print(f"\n=== verification ===")
    print(f"  active_pairs: before={len(before_pairs)}  after={len(after_pairs)}")
    if new_pairs:
        for p in new_pairs:
            print(f"  NEW PAIR: mac={p.get('mac')} game={p.get('name', g['name'])} "
                  f"gsid={(p.get('gsid') or '')[:16]}…")
    else:
        print(f"  no new active_pair for {g['name']} — detector did NOT trigger")
        print(f"  possible causes: source MAC on infra skip list, "
              f"target IP unreachable from this host, or per-flow dedupe "
              f"(if a recent run hit the same MAC+config)")

    journal = gpnc_journal_excerpt(args.gpnc_host, since, "")
    print(f"\n=== gpnc journal (since {since}) ===")
    print(journal or "(no L2-inline launch lines)")
    return 0 if new_pairs else 4


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