#!/usr/bin/env python3
"""Replay the Phew checker's exact interaction, timing every step.

Purpose: find WHICH read exceeds its budget. The receiver log only says
"Timeout waiting for '> '. Got so far: <empty>", which cannot distinguish a
slow keygen from a hang. This prints per-step wall time and the buffer state
at the moment of the timeout.
"""
import os
import select
import subprocess
import sys
import time

CONT = os.environ.get("PHEW_CONT", "phew_container_team1")
CMD = ["docker", "exec", "-i", "-e", "PYTHONUNBUFFERED=1", CONT,
       "python3", "/home/ctfuser/chall/src/chall.py"]


def read_until(proc, needle, timeout):
    """Mirror of the checker's _read_until, but reporting the buffer."""
    buf = ""
    end = time.time() + timeout
    raw = proc.stdout.buffer if hasattr(proc.stdout, "buffer") else proc.stdout
    while needle not in buf:
        left = end - time.time()
        if left <= 0:
            raise TimeoutError(f"timeout after {timeout}s, buffer={buf!r}")
        r, _, _ = select.select([raw], [], [], min(left, 1.0))
        if not r:
            continue
        chunk = raw.read1(4096) if hasattr(raw, "read1") else raw.read(4096)
        if not chunk:
            raise TimeoutError(f"EOF, buffer={buf!r}")
        buf += chunk.decode(errors="replace")
    return buf


def step(label, fn):
    t0 = time.time()
    try:
        out = fn()
        print(f"  {label:<34} {time.time()-t0:6.2f}s  ok")
        return out
    except TimeoutError as e:
        print(f"  {label:<34} {time.time()-t0:6.2f}s  TIMEOUT {e}")
        raise


def main():
    proc = subprocess.Popen(CMD, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                            stderr=subprocess.STDOUT, text=True, bufsize=0)
    try:
        print(f"container={CONT}")
        step("boot -> first menu (budget 45s)", lambda: read_until(proc, "> ", 45))

        for name, send, budget in (("encrypt", "1", 30), ("decrypt", "3", 30),
                                   ("key?", "4", 30)):
            proc.stdin.write(send + "\n")
            proc.stdin.flush()
            step(f"send '{send}' -> prompt (budget 30s)",
                 lambda: read_until(proc, "> ", 30))
            if name == "encrypt":
                step("  send plaintext -> prompt",
                     lambda: read_until(proc, "> ", 30)) if False else None
                proc.stdin.write("414243\n")
                proc.stdin.flush()
                step("  send plaintext -> prompt",
                     lambda: read_until(proc, "> ", 30))
            elif name == "key?":
                proc.stdin.write("\n")
                proc.stdin.flush()
                step("  send blank -> prompt",
                     lambda: read_until(proc, "> ", 30))
            else:
                proc.stdin.write("0\n")
                proc.stdin.flush()
                step("  send ct -> prompt",
                     lambda: read_until(proc, "> ", 30))
        print("ALL STEPS WITHIN BUDGET")
    except TimeoutError:
        print("=> a step needs a bigger budget (or a different prompt)")
        raise SystemExit(1)
    finally:
        try:
            subprocess.run(["docker", "exec", CONT, "pkill", "-f", "chall.py"],
                           capture_output=True, timeout=20)
        except Exception:
            pass
        if proc.poll() is None:
            proc.kill()


if __name__ == "__main__":
    main()