"""Tests for the three relay/console hardening fixes (2026-07-26).

  1. UDP anti-hijack: a datagram may only refresh the endpoint of the host whose
     TCP connection it actually shares an IP with.
  2. Console reassembly: the pod's egg-ACK must be found even when it arrives
     split across recv() calls, or coalesced with another message.
  3. Remote-operator roster adoption: the GUI monitor learns the roster from the
     relay's replayed log line, so a remote operator gets pilot lights and a
     working LAUNCH button.

No sockets, no game, no relay process -- everything is driven through fakes.
"""
import os
import struct
import sys

REPO = r"C:\git\bt411"
sys.path.insert(0, os.path.join(REPO, "tools"))
os.chdir(os.path.join(REPO, "content"))

import btconsole  # noqa: E402

EGG = os.path.join(REPO, "content", "FOGDAY.EGG")
fails = []


def check(label, ok, extra=""):
    print("  %-56s %s %s" % (label, "OK" if ok else "*** FAIL ***", extra))
    if not ok:
        fails.append(label)


def new_relay():
    return btconsole.Relay(1500, EGG, "127.0.0.1", 0, manual_launch=True)


class FakeGameConn:
    """A TCP-registered pod, as by_host holds it."""
    def __init__(self, ip, port=5000, host_id=2):
        self.addr = (ip, port)
        self.host_id = host_id
        self.last_seen = btconsole.time.time()
        self.buf = b""


class FakeUdp:
    """Serves queued datagrams to _udp_read, then blocks like a real socket."""
    def __init__(self, datagrams):
        self.queue = list(datagrams)
        self.sent = []
    def recvfrom(self, _n):
        if not self.queue:
            raise BlockingIOError()
        return self.queue.pop(0)
    def sendto(self, data, endpoint):
        self.sent.append((data, endpoint))


def udp_frame(route, from_host, seq=1, payload=b"\0\0\0\0"):
    return struct.pack(btconsole.ENV_FMT_UDP, route, from_host, seq) + payload


print("=== 1. UDP ANTI-HIJACK ===")
r = new_relay()
r.by_host = {2: FakeGameConn("10.0.0.5", 5000, 2)}
r.udp_endpoint = {2: ("10.0.0.5", 6000)}
r.game_conns = list(r.by_host.values())

# (a) a stranger claiming host 2 must NOT move host 2's endpoint
r.udp_sock = FakeUdp([(udp_frame(btconsole.ROUTE_BCAST, 2), ("66.66.66.66", 7777))])
r._udp_read()
check("spoofed datagram did not move the endpoint",
      r.udp_endpoint[2] == ("10.0.0.5", 6000), str(r.udp_endpoint[2]))
check("spoof counted in stats", r.stats.get("udp_spoofed", 0) == 1)

# (b) the REAL pod on a NEW port (NAT rebind) must still be honoured
r.udp_sock = FakeUdp([(udp_frame(btconsole.ROUTE_BCAST, 2), ("10.0.0.5", 9999))])
r._udp_read()
check("NAT rebind (same IP, new port) accepted",
      r.udp_endpoint[2] == ("10.0.0.5", 9999), str(r.udp_endpoint[2]))

# (c) an unregistered host id is still dropped outright
before = dict(r.udp_endpoint)
r.udp_sock = FakeUdp([(udp_frame(btconsole.ROUTE_BCAST, 99), ("10.0.0.9", 1))])
r._udp_read()
check("unregistered host id dropped", r.udp_endpoint == before)


print("\n=== 2. CONSOLE REASSEMBLY (the egg-ACK must never be missed) ===")

ACK = (struct.pack("<iiii", 0, 0, 2, 0)      # clientID 0, gameID, fromHost, ts
       + struct.pack("<Ii", 12, 4)           # messageLength 12, messageID 4
       + b"\0" * 4)                          # -> 16 + 12 = 28 bytes
assert len(ACK) == 28


class FakeConsoleSock:
    def __init__(self, reads):
        self.reads = list(reads)
    def recv(self, _n):
        return self.reads.pop(0) if self.reads else b""


def acking_relay():
    r = new_relay()
    r.console_conns = []
    r._check_launch_gate = lambda: None      # isolate: we only test detection
    conn = btconsole.RelayConsoleConn(FakeConsoleSock([]), ("10.0.0.5", 5001))
    r.console_conns.append(conn)
    return r, conn


# (a) the ACK arriving in ONE piece (the old code handled only this)
r, conn = acking_relay()
conn.sock.reads = [ACK]
r._console_read(conn)
check("whole ACK detected", conn.acked is True)

# (b) SPLIT across two reads -- the old fixed-offset parse dropped this silently,
#     and a missed ACK means that seat never counts toward the launch gate.
r, conn = acking_relay()
conn.sock.reads = [ACK[:16], ACK[16:]]
r._console_read(conn)
check("split ACK: not detected on the partial half", conn.acked is False)
r._console_read(conn)
check("split ACK detected once the rest arrives", conn.acked is True)

# (c) split at an even nastier boundary (mid-header)
r, conn = acking_relay()
conn.sock.reads = [ACK[:5], ACK[5:20], ACK[20:]]
for _ in range(3):
    r._console_read(conn)
check("ACK split three ways still detected", conn.acked is True)

# (d) two messages COALESCED into one read -- the old code read only the first
r, conn = acking_relay()
other = (struct.pack("<iiii", 0, 0, 2, 0) + struct.pack("<Ii", 12, 99) + b"\0" * 4)
conn.sock.reads = [other + ACK]
r._console_read(conn)
check("coalesced messages: the ACK behind another is found", conn.acked is True)

# (e) an implausible length must drop the connection rather than mis-read forever
r, conn = acking_relay()
dropped = []
r._drop_console = lambda c: dropped.append(c)
conn.sock.reads = [struct.pack("<iiii", 0, 0, 2, 0) + struct.pack("<Ii", 1 << 30, 4)]
r._console_read(conn)
check("garbage messageLength drops the conn", dropped == [conn])


print("\n=== 3. REMOTE OPERATOR: roster adoption + live command channel ===")
try:
    import btoperator  # noqa: E402
except Exception as e:                        # PySide6 missing -> skip, do not fail
    print("  (skipped: cannot import btoperator: %r)" % (e,))
else:
    mon = btoperator.SessionMonitor(True, [])          # how remote mode starts
    check("remote monitor starts with no tags", mon.tags == [])
    line = ("22:27:39 [relay] roster: 2 pilot(s) -> hostIDs [2, 3]: "
            "['10.99.0.1:1502', '10.99.0.2:1502']")
    changed = mon.feed(line)
    check("roster line adopted", mon.tags == ["10.99.0.1:1502", "10.99.0.2:1502"],
          str(mon.tags))
    check("feed reported a change", changed is True)
    check("every tag got a state", len(mon.state) == 2)

    # a seat line now lands on a known tag, so the pilot light works
    mon.feed("[relay] PLAYER 1 SEATED (host 2) tag='10.99.0.1:1502' "
             "callsign='SAURON' mech='lok2'")
    check("SEATED now maps onto an adopted tag",
          mon.state["10.99.0.1:1502"] == btoperator.ST_SEATED)
    seated = sum(1 for s in mon.state.values()
                 if s in (btoperator.ST_SEATED, btoperator.ST_REGISTERED,
                          btoperator.ST_READY))
    check("seated count is no longer stuck at 0", seated == 1, str(seated))

    # re-feeding the same roster must not churn state
    mon.feed(line)
    check("re-feeding the roster preserves known state",
          mon.state["10.99.0.1:1502"] == btoperator.ST_SEATED)

print("\n%s" % ("ALL CHECKS PASSED" if not fails else "FAILURES: %s" % fails))
sys.exit(1 if fails else 0)