graph: docs, and the offline layout tools

cfg_dump freezes real CFGs to JSON; graph_spike renders one or --stats a
whole corpus through the shipping engine; graph_smoke is the end-to-end
tool->domain->layout check; graph_shot renders the real view headless at a
chosen size, which is the only sane way to judge it (a tiled pane is far
too narrow and the minimap sits on top of the graph).

4 files changed, 441 insertions, 0 deletions

diff --git a/experiments/cfg_dump.py b/experiments/cfg_dump.py
new file mode 100644
index 0000000..30b6ba9
--- /dev/null
+++ b/experiments/cfg_dump.py
@@ -0,0 +1,119 @@
+#!/usr/bin/env python3
+"""Dump real basic-block CFGs to JSON, so graph-layout work can be done offline.
+
+Run with a python that has ``idapro`` (i.e. the worker python, /usr/bin/python3):
+
+    /usr/bin/python3 experiments/cfg_dump.py targets/echo -o /tmp/cfg-echo.json
+
+Copies the binary to a scratch dir first (never touches the tracked .i64), opens
+it with idalib, and writes one record per function:
+
+    {"name","ea","blocks":[{"id","start","end","lines":[str],"succs":[[id,kind]]}]}
+
+``kind`` is "fall" (falls through to the next address), "jump" (unconditional or
+taken branch) or "switch" (n-way). That is exactly the input a graph view needs;
+everything after this point (layering, ordering, routing, rendering) is pure
+python and needs no IDA at all.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shutil
+import sys
+import tempfile
+
+
+def dump(path: str, want: list[str], max_blocks: int) -> list[dict]:
+    import idapro
+
+    scratch = tempfile.mkdtemp(prefix="cfgdump-")
+    local = os.path.join(scratch, os.path.basename(path))
+    shutil.copy2(path, local)
+    if idapro.open_database(local, run_auto_analysis=True) != 0:
+        raise SystemExit(f"failed to open {local}")
+
+    import ida_funcs
+    import ida_gdl
+    import ida_lines
+    import ida_bytes
+    import idautils
+    import idaapi
+
+    out = []
+    try:
+        for fea in idautils.Functions():
+            fn = ida_funcs.get_func(fea)
+            if not fn:
+                continue
+            name = ida_funcs.get_func_name(fea)
+            if want and not any(w in name for w in want):
+                continue
+            fc = ida_gdl.FlowChart(fn, flags=ida_gdl.FC_PREDS)
+            blocks = []
+            index = {}
+            for i, bb in enumerate(fc):
+                index[bb.start_ea] = i
+            for bb in fc:
+                lines = []
+                ea = bb.start_ea
+                while ea < bb.end_ea and ea != idaapi.BADADDR:
+                    txt = ida_lines.tag_remove(
+                        ida_lines.generate_disasm_line(ea, 0) or "")
+                    lines.append(txt.rstrip())
+                    nxt = ida_bytes.next_head(ea, bb.end_ea)
+                    if nxt <= ea:
+                        break
+                    ea = nxt
+                succs = []
+                sl = list(bb.succs())
+                for s in sl:
+                    if s.start_ea not in index:
+                        continue
+                    if len(sl) > 2:
+                        kind = "switch"
+                    elif s.start_ea == bb.end_ea:
+                        kind = "fall"
+                    else:
+                        kind = "jump"
+                    succs.append([index[s.start_ea], kind])
+                blocks.append({
+                    "id": index[bb.start_ea],
+                    "start": bb.start_ea,
+                    "end": bb.end_ea,
+                    "lines": lines,
+                    "succs": succs,
+                })
+            if max_blocks and len(blocks) > max_blocks:
+                continue
+            out.append({"name": name, "ea": fea, "blocks": blocks})
+    finally:
+        idapro.close_database(save=False)
+        shutil.rmtree(scratch, ignore_errors=True)
+    return out
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("binary")
+    ap.add_argument("-o", "--out", default="/tmp/cfg.json")
+    ap.add_argument("-f", "--func", action="append", default=[],
+                    help="only functions whose name contains this (repeatable)")
+    ap.add_argument("--max-blocks", type=int, default=0,
+                    help="skip functions with more blocks than this")
+    args = ap.parse_args()
+
+    recs = dump(os.path.abspath(args.binary), args.func, args.max_blocks)
+    recs.sort(key=lambda r: len(r["blocks"]), reverse=True)
+    with open(args.out, "w") as f:
+        json.dump(recs, f)
+    tot = sum(len(r["blocks"]) for r in recs)
+    print(f"{len(recs)} functions, {tot} blocks -> {args.out}")
+    for r in recs[:12]:
+        print(f"  {len(r['blocks']):4d} blocks  {r['name']}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/experiments/graph_shot.py b/experiments/graph_shot.py
new file mode 100644
index 0000000..5f0d45a
--- /dev/null
+++ b/experiments/graph_shot.py
@@ -0,0 +1,66 @@
+"""Render the graph view headless at a chosen size and print it.
+
+    ~/ida-venv/bin/python experiments/graph_shot.py [func] [cols] [rows] [zoom]
+
+The pane a person runs this in is usually too small to judge the layout, and the
+pilot lays out synchronously at whatever size you ask for -- so this is the way
+to actually look at the thing.
+"""
+import asyncio
+import os
+import shutil
+import sys
+
+REPO = os.path.expanduser("~/dev/ida-tui-maybe")
+sys.path.insert(0, REPO)
+os.chdir(REPO)
+
+func = sys.argv[1] if len(sys.argv) > 1 else "sub_2297"
+cols = int(sys.argv[2]) if len(sys.argv) > 2 else 170
+rows = int(sys.argv[3]) if len(sys.argv) > 3 else 55
+zoom = int(sys.argv[4]) if len(sys.argv) > 4 else 0
+
+src, tmp = f"{REPO}/targets/echo", "/tmp/echo_shot"
+shutil.copy(src, tmp)
+for e in ".i64 .id0 .id1 .id2 .nam .til".split():
+    try:
+        os.remove(tmp + e)
+    except OSError:
+        pass
+
+from idatui.app import IdaTui, GraphView   # noqa: E402
+from idatui._sync import wait_for          # noqa: E402
+from idatui.rpc import screen_text          # noqa: E402
+
+
+async def main() -> None:
+    app = IdaTui(tmp, keepalive=False)
+    async with app.run_test(size=(cols, rows)) as pilot:
+        await wait_for(lambda: app.program is not None and app._cur is not None,
+                       pilot.pause, 120)
+        ea = app.program.resolve(func)
+        fn = app.program.function_of(ea)
+        app._open_function(fn.addr, fn.name)
+        await wait_for(lambda: app._cur is not None and app._cur.ea == fn.addr,
+                       pilot.pause, 60)
+        await pilot.pause(0.2)
+        await pilot.press("space")
+        gv = app.query_one(GraphView)
+        ok = await wait_for(lambda: app._active == "graph" and gv.lay is not None,
+                            pilot.pause, 90)
+        if not ok:
+            print("graph never opened:", app.query_one("#status").render())
+            return
+        for _ in range(zoom):
+            await pilot.press("z")
+            await pilot.pause(0.2)
+        await pilot.pause(0.4)
+        print(screen_text(app)["text"])
+        print()
+        print("status:", app.query_one("#status").render())
+        print("stats :", gv.lay.stats, "canvas",
+              f"{gv.lay.width}x{gv.lay.height}", "zoom", gv.ZOOMS[gv._zoom])
+        app._save_on_exit = False
+
+
+asyncio.run(main())
diff --git a/experiments/graph_smoke.py b/experiments/graph_smoke.py
new file mode 100644
index 0000000..1b6e21e
--- /dev/null
+++ b/experiments/graph_smoke.py
@@ -0,0 +1,107 @@
+"""End-to-end smoke for the graph path: the `flowchart` tool -> domain ->
+idatui.graph layout, through a real idalib worker.
+
+    ~/ida-venv/bin/python experiments/graph_smoke.py [funcname]
+
+Wants: VERDICT: OK
+"""
+import os
+import shutil
+import sys
+import time
+
+REPO = os.path.expanduser("~/dev/ida-tui-maybe")
+sys.path.insert(0, REPO)
+os.chdir(REPO)
+
+src, tmp = f"{REPO}/targets/echo", "/tmp/echo_graph"
+shutil.copy(src, tmp)
+for e in ".i64 .id0 .id1 .id2 .nam .til".split():
+    try:
+        os.remove(tmp + e)
+    except OSError:
+        pass
+
+from idatui.worker_client import WorkerClient   # noqa: E402
+from idatui.domain import Program               # noqa: E402
+from idatui import graph as G                   # noqa: E402
+
+want = sys.argv[1] if len(sys.argv) > 1 else "main"
+print("spawning worker + opening echo\u2026", flush=True)
+t = time.time()
+cl = WorkerClient(tmp)
+cl.connect(progress=lambda m: None)
+print(f"  worker ready in {time.time()-t:.2f}s", flush=True)
+prog = Program(cl)
+
+ok = True
+ea = prog.resolve(want)
+print(f"resolve({want!r}) = {ea:#x}", flush=True)
+
+t = time.time()
+fc = prog.flowchart(ea)
+print(f"flowchart() -> {time.time()-t:.2f}s", flush=True)
+if fc is None:
+    print("VERDICT: FAIL (no flowchart)")
+    raise SystemExit(1)
+
+print(f"  {fc.name} @ {fc.func_ea:#x}: {len(fc.blocks)} blocks, entry={fc.entry}")
+nrows = sum(len(b.rows) for b in fc.blocks)
+print(f"  {nrows} listing rows attached")
+ok &= nrows > 0
+if not nrows:
+    print("  !! no rows attached to any block")
+
+empty = [b for b in fc.blocks if not b.rows]
+if empty:
+    ok = False
+    print(f"  !! {len(empty)} blocks have NO rows, e.g. "
+          f"{[hex(b.start) for b in empty[:4]]}")
+
+b0 = fc.blocks[fc.entry]
+print(f"  entry block {b0.start:#x}-{b0.end:#x}:")
+for h in b0.rows[:5]:
+    tags = ",".join(k for k, _ in (h.spans or ()))[:40]
+    print(f"    {h.ea:#010x} {h.kind:<7} {h.text[:42]!r}  spans[{tags}]")
+ok &= any(h.spans for b in fc.blocks for h in b.rows)
+if not any(h.spans for b in fc.blocks for h in b.rows):
+    print("  !! no colour spans came through \u2014 highlighting would be dead")
+
+# rows must tile the block exactly, or boxes will have holes
+for b in fc.blocks:
+    eas = [h.ea for h in b.rows]
+    if eas and (min(eas) < b.start or max(eas) >= b.end):
+        ok = False
+        print(f"  !! block {b.start:#x} has rows outside its range")
+
+blocks = [G.Block(id=b.id, start=b.start, end=b.end, succs=list(b.succs))
+          for b in fc.blocks]
+
+
+def sizer(b):
+    src_b = fc.blocks[b.id]
+    w = max([len(f"loc_{b.start:X}")]
+            + [len(h.text) + 12 for h in src_b.rows]) + 4
+    return (w, len(src_b.rows) + 3)
+
+
+t = time.time()
+lay = G.layout(blocks, sizer, entry=fc.entry)
+print(f"layout() -> {(time.time()-t)*1000:.0f} ms  {lay.stats}", flush=True)
+print(f"  canvas {lay.width}x{lay.height}")
+ok &= len(lay.nodes) == len(blocks)
+
+# every block must be reachable in the drawing, and rows must index it
+covered = {n.id for r in range(lay.height) for n in lay.nodes_at_row(r)}
+ok &= covered == {b.id for b in blocks}
+if covered != {b.id for b in blocks}:
+    print(f"  !! row index misses {sorted({b.id for b in blocks} - covered)[:5]}")
+
+hits = sum(1 for r in range(min(lay.height, 400))
+           if lay.painting.cells_at_row(r, 0, lay.width))
+print(f"  {hits} of the first {min(lay.height,400)} rows carry edge cells")
+ok &= hits > 0
+
+cl.close()
+print("VERDICT:", "OK" if ok else "FAIL")
+raise SystemExit(0 if ok else 1)
diff --git a/experiments/graph_spike.py b/experiments/graph_spike.py
new file mode 100644
index 0000000..547c829
--- /dev/null
+++ b/experiments/graph_spike.py
@@ -0,0 +1,149 @@
+#!/usr/bin/env python3
+"""Render a CFG with idatui.graph, offline, to stdout.
+
+The layout engine now lives in ``idatui/graph.py`` (this started life as the
+spike that proved it). What survives here is the offline harness: feed it a
+corpus from ``cfg_dump.py`` and look at a function, or lay out the whole corpus
+and check the cost. No IDA, no Textual, no worker -- so it iterates in
+milliseconds when you're changing layout heuristics.
+
+    /usr/bin/python3 experiments/cfg_dump.py targets/echo -o /tmp/cfg-echo.json
+    python3 experiments/graph_spike.py /tmp/cfg-echo.json --func sub_61D0
+    python3 experiments/graph_spike.py /tmp/cfg-echo.json --stats
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from idatui import graph as G  # noqa: E402
+
+COLOR = {
+    G.E_UNCOND: "\033[38;5;39m", G.E_TRUE: "\033[38;5;40m",
+    G.E_FALSE: "\033[38;5;203m", G.E_SWITCH: "\033[38;5;178m",
+    G.E_BACK: "\033[38;5;135m",
+}
+DIM, RESET = "\033[38;5;244m", "\033[0m"
+PAD = 1
+
+
+def build(rec: dict, max_lines: int):
+    """(blocks, sizer, texts) for one cfg_dump record."""
+    texts: dict[int, list[str]] = {}
+    for b in rec["blocks"]:
+        lines = list(b["lines"])
+        if max_lines and len(lines) > max_lines:
+            lines = lines[:max_lines - 1] + [f"... {len(b['lines']) - max_lines + 1} more"]
+        texts[b["id"]] = lines
+    blocks = [G.Block(id=b["id"], start=b["start"], end=b["end"],
+                      succs=[(d, k) for d, k in b["succs"]])
+              for b in rec["blocks"]]
+
+    def sizer(b: G.Block) -> tuple[int, int]:
+        lines = texts[b.id]
+        label = f"loc_{b.start:X}"
+        widest = max([len(label) + 4] + [len(l) for l in lines] or [4])
+        return (widest + 2 * PAD + 2, max(len(lines), 1) + 2)
+
+    return blocks, sizer, texts
+
+
+def render(lay: G.Layout, texts: dict[int, list[str]], color: bool) -> str:
+    rows = []
+    for row in range(lay.height):
+        cells: dict[int, tuple[str, str]] = {}
+        for col, (ch, kind, _eid) in lay.painting.cells_at_row(
+                row, 0, lay.width).items():
+            cells[col] = (ch, COLOR.get(kind, ""))
+        for n in lay.nodes_at_row(row):
+            label = f"loc_{n.block.start:X}" if n.block else ""
+            if row == n.y:
+                cells[n.x] = (G.BOX["tl"], DIM)
+                for i in range(1, n.w - 1):
+                    cells[n.x + i] = (G.BOX["h"], DIM)
+                cells[n.x + n.w - 1] = (G.BOX["tr"], DIM)
+                tag = f" {label} "
+                if len(tag) <= n.w - 4:
+                    for k, c in enumerate(tag):
+                        cells[n.x + 2 + k] = (c, DIM)
+                if n.block is not None and n.block.selfloop:
+                    cells[n.x + n.w - 2] = ("\u21ba", COLOR[G.E_BACK])
+            elif row == n.y + n.h - 1:
+                cells[n.x] = (G.BOX["bl"], DIM)
+                for i in range(1, n.w - 1):
+                    cells[n.x + i] = (G.BOX["h"], DIM)
+                cells[n.x + n.w - 1] = (G.BOX["br"], DIM)
+            else:
+                cells[n.x] = (G.BOX["v"], DIM)
+                cells[n.x + n.w - 1] = (G.BOX["v"], DIM)
+                for i in range(1, n.w - 1):
+                    cells[n.x + i] = (" ", "")
+                lines = texts.get(n.id, [])
+                i = row - n.y - 1
+                if 0 <= i < len(lines):
+                    for k, c in enumerate(lines[i]):
+                        cells[n.x + 1 + PAD + k] = (c, "")
+        line, cur, last = [], "", -1
+        for col in sorted(cells):
+            ch, st = cells[col]
+            line.append(" " * (col - last - 1))
+            if color and st != cur:
+                line.append(st or RESET)
+                cur = st
+            line.append(ch)
+            last = col
+        if color and cur:
+            line.append(RESET)
+        rows.append("".join(line))
+    return "\n".join(rows)
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("corpus")
+    ap.add_argument("--func", help="function name (default: the smallest)")
+    ap.add_argument("--stats", action="store_true", help="lay out the whole corpus")
+    ap.add_argument("--max-lines", type=int, default=8,
+                    help="collapse blocks longer than this (0 = never)")
+    ap.add_argument("--no-color", action="store_true")
+    args = ap.parse_args()
+
+    recs = json.load(open(args.corpus))
+    if args.stats:
+        print(f"{'blocks':>7} {'nodes':>6} {'dummy':>6} {'layer':>6} "
+              f"{'canvas':>12} {'ms':>8}  name")
+        tot = 0.0
+        for r in sorted(recs, key=lambda r: len(r["blocks"])):
+            blocks, sizer, _ = build(r, args.max_lines)
+            lay = G.layout(blocks, sizer)
+            s = lay.stats
+            tot += s["ms"]
+            print(f"{s['blocks']:>7} {s['nodes']:>6} {s['dummies']:>6} "
+                  f"{s['layers']:>6} {lay.width:>5}x{lay.height:<6} "
+                  f"{s['ms']:>8.1f}  {r['name']}")
+        print(f"total {tot:.0f} ms over {len(recs)} functions")
+        return 0
+
+    if args.func:
+        rec = next((r for r in recs if r["name"] == args.func), None)
+        if rec is None:
+            print("no such function; have: "
+                  f"{', '.join(r['name'] for r in recs[:20])}", file=sys.stderr)
+            return 1
+    else:
+        rec = min(recs, key=lambda r: len(r["blocks"]))
+
+    blocks, sizer, texts = build(rec, args.max_lines)
+    lay = G.layout(blocks, sizer)
+    print(render(lay, texts, color=not args.no_color))
+    print(f"\n{rec['name']}: {lay.stats}  canvas {lay.width}x{lay.height}",
+          file=sys.stderr)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())