Syplheed-Reborn/tools/re-capture/ob_flag.py

#!/usr/bin/env python3
"""Does an `OB`-badged entity carry a flag, and is `REMAINING OB` its count?

`mission-objective-counter.md` has the counter's address; what it *counts* is the
part the autopilot needs, because "shoot what closes the mission" requires
picking the right target, not knowing how many are left.

Two questions, in order of how cheaply they can be killed:

1. **Is the counter just a per-class head-count?** Print the class histogram
   beside the counter. The corpus already suspects not (012 on the HUD against
   118 live ADAN), and one run settles it for every class at once.
2. **Is there a per-entity flag whose set-cardinality is the counter?** For every
   4-byte offset in a window around each entity, count how many entities share
   each value. An offset where exactly N entities agree, with N the counter, is a
   candidate — and there will be many by chance, so the answer is the SECOND
   sample: after the counter moves to N', the same (offset, value) must be shared
   by exactly N' entities. That is the same "verify across a transition you did
   not select on" rule the address itself had to pass.

🔴 Known limit, stated because it bounds the conclusion: `entities2.typed`
enumerates entities by their position triple CHANGING between two samples, so a
stationary objective is invisible to it. Stage 02's objective is "shoot down all
invading enemy fighters", which move — but a null result here does not rule out a
flag on objects this enumeration never sees.

Usage: ob_flag.py <out-dir> [timeout_s]
"""
import json
import os
import struct
import subprocess
import sys
import time
from collections import Counter, defaultdict

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import entities2  # noqa: E402
import frozen     # noqa: E402
import gmem       # noqa: E402
import gworld     # noqa: E402
import ob_read    # noqa: E402

VA = 0xBDB59668
RADIUS = 0x400
DELTA = 0x130


def counter(fd):
    return struct.unpack(">I", os.pread(fd, 4, gmem.va_to_off(VA)))[0]


def hud(shot):
    subprocess.run(["screenshot", shot], stdout=subprocess.DEVNULL,
                   stderr=subprocess.DEVNULL)
    txt, _ = ob_read.read(shot)
    return int(txt) if txt.isdigit() else None


def sample(w, defs, out, tag):
    """(counter, entity list, {(offset, value): [entity positions]})."""
    fd = w.fd
    n = counter(fd)
    movers = entities2.moving(fd, w.size)
    ents = entities2.typed(fd, defs, movers, DELTA)
    uniq = {}
    for off, nm, pos, sp in ents:
        uniq.setdefault((nm, tuple(round(c, 1) for c in pos)), (off, nm))
    ents = list(uniq.values())
    groups = defaultdict(list)
    for off, nm in ents:
        lo = off - RADIUS
        blob = os.pread(fd, RADIUS * 2, lo)
        for k in range(0, len(blob) - 3, 4):
            groups[(lo + k - off, blob[k:k + 4])].append(nm)
    print(f"[{tag}] counter={n} entities={len(ents)}", flush=True)
    return n, ents, groups


def main():
    out = sys.argv[1]
    deadline = time.time() + (float(sys.argv[2]) if len(sys.argv) > 2 else 600)
    os.makedirs(out, exist_ok=True)
    w = gworld.World()
    defs = entities2.definitions(w)

    v = hud(f"{out}/a.png")
    n0 = counter(w.fd)
    print(f"HUD={v} RAM={n0}", flush=True)
    if v != n0:
        print("HUD and RAM disagree — wrong address for this run; re-scan with "
              "ob_hunt.py before trusting anything below", flush=True)
        return 2

    nA, entsA, gA = sample(w, defs, out, "A")
    hist = Counter(nm for _, nm in entsA)
    print(f"[A] class histogram vs counter {nA}:", flush=True)
    for nm, k in hist.most_common(12):
        print(f"      {k:4d}  {nm}", flush=True)
    exact = [nm for nm, k in hist.items() if k == nA]
    print(f"[A] classes whose head-count equals the counter: {exact or 'NONE'}",
          flush=True)

    candA = {k: v for k, v in gA.items() if len(v) == nA}
    print(f"[A] offsets where exactly {nA} entities agree: {len(candA)}", flush=True)

    stuck = 0
    while time.time() < deadline:
        time.sleep(5)
        if counter(w.fd) != nA:
            break
        stuck += 1
        if stuck % 12 == 0 and frozen.frozen(6.0)[0]:
            print("GUEST FROZEN — the world stopped advancing, so the counter "
                  "was never going to move; this run proves nothing", flush=True)
            return 3
    nB = counter(w.fd)
    if nB == nA:
        print("counter never moved — no verification possible", flush=True)
        return 1
    vb = hud(f"{out}/b.png")
    print(f"counter {nA} -> {nB} (HUD {vb})", flush=True)

    _, entsB, gB = sample(w, defs, out, "B")
    survivors = {k: (len(gA[k]), len(gB.get(k, []))) for k in candA
                 if len(gB.get(k, [])) == nB}
    print(f"[B] of {len(candA)} candidates, {len(survivors)} still hold "
          f"exactly {nB}", flush=True)
    rows = [{"delta": d, "value": val.hex(), "a": a, "b": b}
            for (d, val), (a, b) in sorted(survivors.items())]
    json.dump({"nA": nA, "nB": nB, "hud_a": v, "hud_b": vb,
               "entities_a": len(entsA), "entities_b": len(entsB),
               "class_matches": exact, "candidates_a": len(candA),
               "survivors": rows}, open(f"{out}/flag.json", "w"), indent=1)
    for r in rows[:40]:
        print(f"   pos{r['delta']:+#07x} = {r['value']}  {r['a']} -> {r['b']}",
              flush=True)
    print(f"wrote {out}/flag.json", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())