#!/usr/bin/env python3
"""Identify shared runtime routines across the TAs and propagate one name to all of them.

These twelve TAs are built against the same PSP TEE runtime, so the same helper appears
byte-identical in many of them. Naming such a routine once and propagating it by content
hash is far better leverage than renaming 384 anonymous functions one at a time, and it is
verifiable: a cluster is only accepted when the bytes match exactly.

Two modes:
  dump    --prefix X --body B     : record IDA's function list + per-function content hash
  report                          : cluster across all dumps, rank by how many TAs share it
  apply   --prefix X              : rename that TA's members of every named cluster

Cluster names live in fw/libnames.json ({hash: name}), filled in from inspection, so the
mapping is auditable rather than implicit.

usage: IDATUI_RPC_SOCK=... ta_libmatch.py dump --prefix gcp --body <body.bin>
       ta_libmatch.py report
       IDATUI_RPC_SOCK=... ta_libmatch.py apply --prefix gcp
"""
import argparse
import glob
import hashlib
import json
import os
import sys
import tempfile

CWD = os.path.expanduser('~/dev/ida-tui-maybe')
sys.path.insert(0, CWD)
from idatui.rpcclient import RpcClient                                   # noqa: E402

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DUMPS = os.path.join(REPO, 'fw', 'libdumps')
NAMES = os.path.join(REPO, 'fw', 'libnames.json')
BASE = 0x400100


def connect():
    return RpcClient(os.environ['IDATUI_RPC_SOCK'], timeout=60)


def cmd_dump(a):
    c = connect()
    c.connect()
    fs = c.call('functions', filter='', limit=2000)
    body = open(a.body, 'rb').read()
    out = []
    for f in fs:
        ea, sz, nm = f.get('ea'), f.get('size') or 0, f.get('name')
        off = ea - BASE
        if off < 0 or sz <= 0 or off + sz > len(body):
            continue
        h = hashlib.sha256(body[off:off + sz]).hexdigest()
        out.append(dict(ea=ea, size=sz, name=nm, hash=h))
    os.makedirs(DUMPS, exist_ok=True)
    json.dump(out, open(os.path.join(DUMPS, '%s.json' % a.prefix), 'w'), indent=1)
    print('%s: %d functions dumped (%d hashable)' % (a.prefix, len(fs), len(out)))
    return 0


def load_dumps():
    d = {}
    for p in glob.glob(os.path.join(DUMPS, '*.json')):
        d[os.path.basename(p)[:-5]] = json.load(open(p))
    return d


def cmd_report(a):
    dumps = load_dumps()
    clusters = {}
    for ta, fns in dumps.items():
        for f in fns:
            clusters.setdefault(f['hash'], []).append((ta, f))
    shared = {h: v for h, v in clusters.items() if len({x[0] for x in v}) >= 2}
    print('TAs dumped: %d ; functions: %d ; distinct bodies: %d ; shared across >=2 TAs: %d'
          % (len(dumps), sum(len(v) for v in dumps.values()), len(clusters), len(shared)))
    named = json.load(open(NAMES)) if os.path.exists(NAMES) else {}
    rows = sorted(shared.items(), key=lambda kv: -len({x[0] for x in kv[1]}))
    print('\n%-12s %-5s %-5s %-8s %s' % ('hash', 'TAs', 'size', 'named?', 'example names'))
    for h, v in rows[:40]:
        tas = sorted({x[0] for x in v})
        ex = sorted({x[1]['name'] for x in v})[:3]
        print('%-12s %-5d %-5d %-8s %s'
              % (h[:10], len(tas), v[0][1]['size'],
                 named.get(h, '-')[:8], ', '.join(ex)[:58]))
    unnamed = [h for h, v in rows if h not in named]
    print('\nshared clusters still unnamed: %d (covering %d function instances)'
          % (len(unnamed), sum(len(clusters[h]) for h in unnamed)))
    return 0


def cmd_apply(a):
    names = json.load(open(NAMES))
    dumps = load_dumps()
    fns = dumps.get(a.prefix, [])
    syms = [dict(addr=f['ea'], name='%s_%s' % (a.prefix, names[f['hash']]))
            for f in fns if f['hash'] in names and f['name'].startswith('sub_')]
    if not syms:
        print('%s: nothing to apply' % a.prefix)
        return 0
    c = connect()
    c.connect()
    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('%s: applied %s/%s shared-library names' % (a.prefix, m.get('ok'),
                                                      m.get('requested')))
    return 0


def main():
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest='mode', required=True)
    d = sub.add_parser('dump')
    d.add_argument('--prefix', required=True)
    d.add_argument('--body', required=True)
    d.set_defaults(fn=cmd_dump)
    r = sub.add_parser('report')
    r.set_defaults(fn=cmd_report)
    p = sub.add_parser('apply')
    p.add_argument('--prefix', required=True)
    p.set_defaults(fn=cmd_apply)
    a = ap.parse_args()
    return a.fn(a)


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