#!/usr/bin/env python3
"""Attribute the remaining anonymous functions to the named function that reaches them.

After shared-runtime propagation, what is left in each TA is TA-specific internal logic.
Renaming it one function at a time invents nothing but also records nothing. What *is*
derivable is **who uses it**: if exactly one named function calls an anonymous one, the
helper belongs to that command/handler, and saying so is a fact about the call graph.

So this names such functions `<caller>__hN` and iterates to a fixpoint, so chains of
helpers behind one command all get attributed to it. Functions with several distinct named
callers are left alone (shared internals -- claiming one owner would be wrong), as are
functions with no named caller.

The call graph is built locally from IDA's own function boundaries (`fw/libdumps/*.json`)
plus a capstone decode of each function's bytes, so it costs no RPC round-trips and cannot
drift from IDA's idea of where functions start.

usage: ta_attribute.py --prefix gcp --body <body.bin> [--apply]   (IDATUI_RPC_SOCK for apply)
"""
import argparse
import json
import os
import sys
import tempfile

import capstone as cs

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


def call_graph(body, fns):
    """{caller_ea: set(callee_ea)} using IDA's boundaries as decode anchors."""
    starts = {f['ea'] for f in fns}
    g = {}
    for f in fns:
        off = f['ea'] - BASE
        out = set()
        for i in MD.disasm(body[off:off + 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)
        g[f['ea']] = out
    return g


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()
    g = call_graph(body, fns)
    name = {f['ea']: f['name'] for f in fns}
    anon = {ea for ea, n in name.items() if n.startswith('sub_')}

    # invert: callee -> set of callers
    callers = {}
    for src, dsts in g.items():
        for d in dsts:
            callers.setdefault(d, set()).add(src)

    assigned, rounds = {}, 0
    while True:
        rounds += 1
        progress = False
        for ea in sorted(anon):
            if ea in assigned:
                continue
            named_callers = {name[c] if c not in assigned else assigned[c]
                             for c in callers.get(ea, set())
                             if (c not in anon) or (c in assigned)}
            # strip our own generated suffix so a chain attributes to the root owner
            roots = {n.split('__h')[0] for n in named_callers}
            if len(roots) == 1:
                root = roots.pop()
                n = sum(1 for v in assigned.values() if v.split('__h')[0] == root) + 1
                assigned[ea] = '%s__h%d' % (root, n)
                progress = True
        if not progress or rounds > 12:
            break

    multi = sum(1 for ea in anon if ea not in assigned and
                len({name[c].split('__h')[0] for c in callers.get(ea, set())
                     if c not in anon}) > 1)
    orphan = sum(1 for ea in anon if ea not in assigned and not callers.get(ea))
    print('%s: %d anonymous -> %d attributed (%d rounds); %d have multiple named owners, '
          '%d have no caller'
          % (a.prefix, len(anon), len(assigned), rounds, multi, orphan))

    if a.apply and assigned:
        from idatui.rpcclient import RpcClient
        c = RpcClient(os.environ['IDATUI_RPC_SOCK'], timeout=60)
        c.connect()
        syms = [dict(addr=ea, name=nm) for ea, nm in assigned.items()]
        with tempfile.NamedTemporaryFile('w', suffix='.json', delete=False) as fh:
            json.dump(syms, fh)
            p = fh.name
        r = c.call('rename_many', file=p) or {}
        m = r.get('rename_many', r)
        os.unlink(p)
        c.call('save')
        print('   applied %s/%s' % (m.get('ok'), m.get('requested')))
    return 0


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