#!/usr/bin/env python3 """Propagate function names across builds through the CALL GRAPH, not through byte similarity. Byte and mnemonic porting names a function only when its own body still resembles the reference. Everything it cannot reach stays sub_*, and for ASD that is thousands of functions. But a function's position in the call graph is evidence independent of its body: if reference function A is matched to target function A', and both make the same number of calls, then the Nth callee of A' is the Nth callee of A. That is the rule here, with three deliberate restrictions, because positional alignment is exactly the kind of heuristic that quietly produces wrong names: * EQUAL CALL COUNT. If A and A' do not make the same number of calls, the alignment is not trustworthy and the anchor is skipped entirely rather than aligned partially. * UNIQUE CANDIDATE. A target function is named only if every anchor that reaches it agrees on the name. Any disagreement and it is left alone. * SIZE AGREEMENT. The candidate and its reference must be within a factor of `--size-ratio` of each other. A name is an identity claim, and a function half the size of the one whose name it is being given is probably not that function. Names are applied to the target database, so re-running fw/ta_re.py afterwards picks them up and re-emits the artifacts. Rounds are iterated: a function named in round N becomes an anchor in round N+1. Control: --dry-run reports what would be applied without touching anything, and the ANCHOR-MISMATCH rate is printed, because a high mismatch rate means the builds have diverged too far for positional alignment to mean anything. usage: ./ta_cgname.py --type asd [--rounds 3] [--dry-run] """ from __future__ import annotations import argparse import glob import json import os import re import sys from pathlib import Path HERE = Path(__file__).resolve().parent NAMEDIR_REL = "cgnames" REPO = HERE.parent sys.path.insert(0, str(Path.home() / "dev" / "ida-tui-maybe")) sys.path.insert(0, str(HERE)) CALL = re.compile(r"\b([A-Za-z_]\w*)\s*\(") KEYWORDS = {"if", "while", "for", "switch", "return", "sizeof", "__fastcall", "do", "else", "case", "break", "continue", "goto"} SUBNAME = re.compile(r"^sub_[0-9A-Fa-f]+$") def func_calls(client, remote_ops, addr): r = client.call(remote_ops.decompile, addr=addr, include_addresses=False) or {} code = r.get("code") or "" code = "\n".join(l for l in code.splitlines() if not l.lstrip().startswith("//")) body = code[code.find("{"):] if "{" in code else "" return [m.group(1) for m in CALL.finditer(body) if m.group(1) not in KEYWORDS] def list_funcs(client, remote_ops): out, off = [], 0 while True: r = client.call(remote_ops.list_funcs, queries=[{"count": 500, "offset": off}]) blk = r["result"][0] out.extend(blk["data"]) off = blk["next_offset"] if off >= blk["total"]: break return out def main() -> int: ap = argparse.ArgumentParser() ap.add_argument("--type", required=True) ap.add_argument("--rounds", type=int, default=3) ap.add_argument("--size-ratio", type=float, default=2.0) ap.add_argument("--dry-run", action="store_true") ap.add_argument("--ref", help="use this build as the reference instead of the registered " "one (substring match); lets a well-named CORPUS build seed " "its weaker siblings, which the handover reference cannot " "always reach") a = ap.parse_args() from idatui import remote_ops from idatui.nexus_client import NexusClient import ta_re as T if a.ref: cands = [p for p in sorted(glob.glob(str(T.ELF_DIR / f"{a.type}_*.elf"))) if a.ref in os.path.basename(p)] if not cands: raise SystemExit("no build matching --ref %r" % a.ref) refelf = cands[0] elif a.type in T.HANDOVER_REF: refelf = str(REPO / "psp_out" / "handover" / "elf" / f"{a.type}2.elf") else: ver = T.LIBDUMP_REF[a.type] refelf = sorted(glob.glob(str(T.ELF_DIR / f"{a.type}_{ver}_*.elf")))[0] targets = [p for p in sorted(glob.glob(str(T.ELF_DIR / f"{a.type}_*.elf"))) if p != refelf] if not targets: print("no target builds for %s" % a.type) return 0 # Names are persisted to fw/cgnames/.json as well as written into the database. # Writing only into the .i64 has already cost this work twice: three databases have come # back corrupted, and rebuilding one discards every name derived here (GCP fell from 83% # to 74% that way). The JSON is the durable record; fw/ta_re.py re-applies it. namedir = HERE / NAMEDIR_REL namedir.mkdir(exist_ok=True) mapfile = namedir / f"{a.type}.json" persisted = json.loads(mapfile.read_text()) if mapfile.exists() else {} rc = NexusClient(refelf) rc.connect() try: rf = list_funcs(rc, remote_ops) rsize = {f["name"]: f["size"] for f in rf} raddr = {f["name"]: f["addr"] for f in rf if not SUBNAME.match(f["name"]) and not f["name"].startswith("j_")} rcalls: dict[str, list] = {} print("%s: reference %s, %d named anchors available" % (a.type, os.path.basename(refelf), len(raddr)), flush=True) total_applied = 0 for tp in targets: key = os.path.basename(tp).replace(".elf", "") tc = NexusClient(tp) try: tc.connect() applied_here = 0 for rnd in range(a.rounds): tf = list_funcs(tc, remote_ops) tname = {f["addr"]: f["name"] for f in tf} tsize = {f["addr"]: f["size"] for f in tf} byname = {f["name"]: f["addr"] for f in tf} anchors = [n for n in raddr if n in byname] cand: dict[str, set] = {} eq = mismatch = 0 for n in anchors: if n not in rcalls: try: rcalls[n] = func_calls(rc, remote_ops, raddr[n]) except Exception: rcalls[n] = [] rl = rcalls[n] if not rl: continue try: tl = func_calls(tc, remote_ops, byname[n]) except Exception: continue if len(tl) != len(rl): mismatch += 1 continue eq += 1 for x, y in zip(rl, tl): if SUBNAME.match(y) and not SUBNAME.match(x) \ and not x.startswith("j_"): cand.setdefault(y, set()).add(x) # unique candidate + size agreement renames = [] for tgtnm, names in cand.items(): if len(names) != 1: continue new = next(iter(names)) ad_existing = byname.get(new) if ad_existing is not None: # Already named here -- most likely by an earlier run of this very # tool. Record it in the durable map anyway, otherwise a module that # was processed before the map existed can never be recovered from # source and stays hostage to its .i64. persisted.setdefault(key, {})[ad_existing] = new continue ad = byname.get(tgtnm) if ad is None: continue rs, ts = rsize.get(new, 0), tsize.get(ad, 0) if rs and ts: hi, lo = max(rs, ts), min(rs, ts) if lo == 0 or hi / lo > a.size_ratio: continue renames.append({"addr": ad, "name": new}) if not renames: print(" %-30s round %d: anchors eq=%d mismatch=%d -> nothing" % (os.path.basename(tp)[:30], rnd + 1, eq, mismatch), flush=True) break if a.dry_run: print(" %-30s round %d: anchors eq=%d mismatch=%d -> WOULD apply %d" % (os.path.basename(tp)[:30], rnd + 1, eq, mismatch, len(renames)), flush=True) break rr = tc.call(remote_ops.rename, batch={"func": renames, "allow_overwrite": False}) # response shape is {"func":[...], "summary":{"ok":N,"failed":M}} ok = int((rr.get("summary") or {}).get("ok", 0)) applied_here += ok for r_ in renames: persisted.setdefault(key, {})[r_["addr"]] = r_["name"] print(" %-30s round %d: anchors eq=%d mismatch=%d -> applied %d" % (os.path.basename(tp)[:30], rnd + 1, eq, mismatch, ok), flush=True) if ok == 0: break if applied_here and not a.dry_run: try: tc.save_database() except Exception: pass total_applied += applied_here finally: try: tc.close() except Exception: pass if not a.dry_run: mapfile.write_text(json.dumps(persisted, indent=1, sort_keys=True)) print("persisted %d build maps to %s" % (len(persisted), mapfile.relative_to(REPO))) print("%s: %d names applied across %d builds" % (a.type, total_applied, len(targets))) finally: try: rc.close() except Exception: pass return 0 if __name__ == "__main__": sys.exit(main())