import io
import os

os.chdir(r"C:/Users/decid/Documents/projects/spt-core/.worktrees/hertz-inbound-probe")
P = "crates/spt-daemon/tests/twohost_web.rs"
s = io.open(P, encoding="utf-8", newline="").read()

# This file is CRLF in the working copy; a multi-line pattern spelled with LF
# matches nothing and the assert fires. Normalize every pattern to the file's own
# terminator instead of guessing which layer I am editing.
CR, LF = chr(13), chr(10)
CRLF = s.count(CR + LF) > s.count(LF) // 2


def term(t):
    return t.replace(CR + LF, LF).replace(LF, CR + LF) if CRLF else t.replace(CR + LF, LF)

# ── A: the ack leaves the ephemeral socket too. The listening port NEVER sends.
old_ack = """                if let Err(e) = socket.send_to(PROBE_ACK, from) {"""
new_ack = """                // OUT OF THE BEACON SOCKET, not this one: an ack sent from the
                // listening port would write return state for the very tuple the
                // next run must find cold, so the guard would poison its own
                // next measurement at the end of every run, green or red.
                if let Err(e) = beacon_socket.send_to(PROBE_ACK, from) {"""
old_ack, new_ack = term(old_ack), term(new_ack)
assert s.count(old_ack) == 1
s = s.replace(old_ack, new_ack)

# ── B: two sockets. Fixed one RECEIVES beacons; ephemeral one SENDS probes and
#    takes the ack, so no run repeats a 4-tuple that a previous run opened.
old_bind = """    let socket = UdpSocket::bind(SocketAddr::new(IpAddr::from([0, 0, 0, 0]), listen))
        .unwrap_or_else(|e| panic!("INBOUND_PROBE: role B could not bind udp {listen}: {e}"));
    socket
        .set_read_timeout(Some(Duration::from_millis(250)))
        .expect("probe socket takes a read timeout");"""
new_bind = """    let socket = UdpSocket::bind(SocketAddr::new(IpAddr::from([0, 0, 0, 0]), listen))
        .unwrap_or_else(|e| panic!("INBOUND_PROBE: role B could not bind udp {listen}: {e}"));
    socket
        .set_read_timeout(Some(Duration::from_millis(250)))
        .expect("probe socket takes a read timeout");
    // ⚠ THE PROBES LEAVE A SEPARATE, EPHEMERAL SOCKET — this fixed one only ever
    // RECEIVES beacons.
    //
    // With a fixed source port here, the tuple (B:port_b+9 → A:probe) REPEATS
    // from run to run, so any earlier outbound from A to that port leaves return
    // state that carries the next run's probe across a link that is actually
    // blocked — at either layer, since WFP and Tailscale's filter both track
    // outbound UDP flows by 4-tuple. An ephemeral source port makes each run's
    // tuple new, so no run can poison the next BY CONSTRUCTION rather than by a
    // cold-start convention nobody will remember. The ack comes back to this
    // ephemeral port because A replies to `from`; that direction is open anyway
    // and nothing is inferred from it.
    let probe_socket = UdpSocket::bind(SocketAddr::new(IpAddr::from([0, 0, 0, 0]), 0))
        .expect("probe sender binds an ephemeral port");
    probe_socket
        .set_read_timeout(Some(Duration::from_millis(250)))
        .expect("probe sender takes a read timeout");"""
old_bind, new_bind = term(old_bind), term(new_bind)
assert s.count(old_bind) == 1
s = s.replace(old_bind, new_bind)

old_send = """        match socket.send_to(PROBE_MAGIC, target) {"""
new_send = """        match probe_socket.send_to(PROBE_MAGIC, target) {"""
old_send, new_send = term(old_send), term(new_send)
assert s.count(old_send) == 1
s = s.replace(old_send, new_send)

# the ack now arrives on the ephemeral socket; the beacon check that used to
# guard against reading a beacon as an ack stays, because the fixed socket is
# still live and a foreign datagram can land on either.
old_recv = """        match socket.recv_from(&mut buf) {
            Ok((n, from)) if buf[..n] == *PROBE_ACK => {"""
new_recv = """        match probe_socket.recv_from(&mut buf) {
            Ok((n, from)) if buf[..n] == *PROBE_ACK => {"""
old_recv, new_recv = term(old_recv), term(new_recv)
assert s.count(old_recv) == 1
s = s.replace(old_recv, new_recv)

io.open(P, "w", encoding="utf-8", newline="").write(s)
print("socket discipline applied: A listens-only, B sends from an ephemeral port")