#!/usr/bin/env python3
"""Document every still-anonymous function with facts, instead of inventing a name.

After name propagation, call-graph attribution and pointer-table recovery, ~14% of
functions remain `sub_*`. They are in two classes, both established
(`fw/AUDIT-NEW-TAS-RE.md`): helpers with *several* named owners, where picking one owner
would misattribute; and functions reached only indirectly or not at all.

Naming those would be invention. Documenting them is not: for each one this writes an IDA
function comment stating its size, its callers, its callees (named where known), and which
class it falls into. That makes the database self-describing for the residue too — a reader
lands on `sub_401ADA` and immediately sees "no named caller; calls X and Y; island member"
rather than nothing.

Facts come from IDA's own function boundaries (`fw/libdumps/*.json`) plus a capstone decode,
so nothing here depends on a guess.

usage: IDATUI_RPC_SOCK=... ta_document_residue.py --prefix gcp --body <body.bin> [--apply]
"""
import argparse
import json
import os
import sys

import capstone as cs

BASE = 0x400100
REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
MD = cs.Cs(cs.CS_ARCH_ARM, cs.CS_MODE_THUMB)
CWD = os.path.expanduser('~/dev/ida-tui-maybe')
sys.path.insert(0, CWD)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--prefix', required=True)
    ap.add_argument('--body', required=True)
    ap.add_argument('--apply', action='store_true')
    a = ap.parse_args()

    fns = json.load(open(os.path.join(REPO, 'fw', 'libdumps', '%s.json' % a.prefix)))
    body = open(a.body, 'rb').read()
    name = {f['ea']: f['name'] for f in fns}
    starts = set(name)

    callees, callers = {}, {}
    for f in fns:
        out = set()
        for i in MD.disasm(body[f['ea'] - BASE:f['ea'] - BASE + f['size']], f['ea']):
            if i.mnemonic in ('bl', 'blx') and i.op_str.startswith('#'):
                try:
                    t = int(i.op_str[1:], 0)
                except ValueError:
                    continue
                if t in starts:
                    out.add(t)
                    callers.setdefault(t, set()).add(f['ea'])
        callees[f['ea']] = out

    anon = [f for f in fns if f['name'].startswith('sub_')]
    # data-referenced?
    import struct
    words = set()
    for off in range(0, len(body) - 4, 2):
        words.add(struct.unpack_from('<I', body, off)[0])

    notes = []
    for f in anon:
        ea = f['ea']
        cs_ = sorted(callers.get(ea, set()))
        named_owners = {name[c].split('__h')[0] for c in cs_
                        if not name[c].startswith('sub_')}
        cls = ('helper with several named owners (attribution would misattribute)'
               if len(named_owners) > 1 else
               'no named caller; not referenced as data -> indirect or dead'
               if not cs_ and (ea | 1) not in words and ea not in words else
               'no bl caller, but referenced as data -> reached indirectly'
               if not cs_ else
               'called only from still-anonymous code')
        lines = [
            'Residue function, documented rather than named (naming would be invention).',
            '',
            'size: %d bytes' % f['size'],
            'class: %s' % cls,
            'callers: %s' % (', '.join(name[c] for c in cs_) or 'none found'),
            'callees: %s' % (', '.join(sorted(name[t] for t in callees[ea])) or 'none'),
            '',
            'Basis: IDA function boundaries + capstone decode of this function only;',
            'see fw/AUDIT-NEW-TAS-RE.md for how the residue was classified.',
        ]
        notes.append((ea, '\\n'.join(lines)))

    print('%s: %d anonymous functions to document' % (a.prefix, len(notes)))
    for ea, _ in notes[:4]:
        print('   %#x %s' % (ea, name[ea]))

    if a.apply and notes:
        # Headless via the Nexus client. The TUI RPC path (goto + comment) depends on the
        # cursor landing on the right token and cannot be verified without scraping the
        # screen; set_comments takes the address directly and reports per-item status.
        import sys as _sys
        _sys.path.insert(0, os.path.expanduser('~/dev/ida-tui-maybe'))
        from idatui import remote_ops
        from idatui.nexus_client import NexusClient
        elf = json.load(open(os.path.join(REPO, 'fw', 'syms', 'index.json')))[a.prefix]['elf']
        c = NexusClient(os.path.join(REPO, elf))
        c.connect(timeout=600)
        items = [{'addr': hex(ea), 'comment': txt.replace('\\n', '\n')} for ea, txt in notes]
        r = c.call(remote_ops.set_comments, items=items)
        rows = r.get('result', []) if isinstance(r, dict) else []
        ok = sum(1 for x in rows if not x.get('error'))
        print('   documented %d/%d' % (ok, len(notes)))
        c.call_save = getattr(c, 'save_database', None)
        if c.call_save:
            print('   save: %s' % json.dumps(c.save_database())[:60])
        c.close()
    return 0


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