#!/usr/bin/env python3
"""Define the functions IDA missed, found via call sites outside its function set.

Discovered while classifying the annotation residue: GCP has `bl` instructions at
`0x402540..0x40259a` that call *known* functions yet lie outside every IDA function --
they are inside `sub_402524`, the MAC driver analysed by hand earlier, which IDA never
turned into a function. So coverage percentages computed over IDA's function list used an
incomplete denominator, and every detector that iterated IDA functions skipped this code.

Method: sweep the image; wherever a `bl`/`blx` to a known function start sits outside all
IDA functions, walk **back** to the nearest plausible prologue (`push{...,lr}` or
`sub sp,#imm`) and define a function there. Re-dumping afterwards shows the new count.

Deliberately conservative: only regions that *call a known function* are considered, which
is a strong code signal. A permissive "does it decode as Thumb" test was tried first and
rejected -- it claimed 75% of these images were unanalysed code, when most of that is data
(OPM embeds XML certificates). That is the same fallacy as judging an ISA by decode
percentage, which this project already recorded.

usage: IDATUI_RPC_SOCK=... ta_define_missed.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)))
    b = open(a.body, 'rb').read()
    starts = {f['ea'] for f in fns}
    inside = set()
    for f in fns:
        inside.update(range(f['ea'], f['ea'] + f['size']))

    # 1. unmarked call sites into known functions
    sites, off = [], 0x401000 - BASE
    while off < len(b) - 2:
        addr = BASE + off
        got = None
        for i in MD.disasm(b[off:off + 4], addr):
            got = i
            break
        if not got:
            off += 2
            continue
        if addr not in inside and got.mnemonic in ('bl', 'blx') and \
                got.op_str.startswith('#'):
            try:
                t = int(got.op_str[1:], 0)
            except ValueError:
                t = None
            if t in starts:
                sites.append(addr)
        off += got.size

    # 2. walk back from each site to a prologue
    cands = set()
    for s in sites:
        for back in range(2, 0x300, 2):
            a0 = s - back
            if a0 in inside or a0 < 0x401000:
                break
            got = None
            for i in MD.disasm(b[a0 - BASE:a0 - BASE + 4], a0):
                got = i
                break
            if not got:
                continue
            is_push = got.mnemonic.startswith('push') and 'lr' in got.op_str
            is_sub = (got.mnemonic in ('sub', 'sub.w') and 'sp' in got.op_str
                      and '#' in got.op_str)
            if is_push or is_sub:
                # a `sub sp` is usually the SECOND instruction of the prologue; prefer the
                # push in front of it, or the function starts mid-prologue (GCP's real
                # start is 0x402524 `push`, not 0x402528 `sub sp`).
                if is_sub:
                    for delta in (2, 4):
                        pa = a0 - delta
                        g2 = None
                        for i2 in MD.disasm(b[pa - BASE:pa - BASE + 4], pa):
                            g2 = i2
                            break
                        if g2 and g2.mnemonic.startswith('push') and 'lr' in g2.op_str:
                            a0 = pa
                            break
                cands.add(a0)
                break
    print('%s: %d unmarked call sites -> %d candidate prologues %s'
          % (a.prefix, len(sites), len(cands),
             sorted(hex(x) for x in cands)[:8]))

    if a.apply and cands:
        from idatui.rpcclient import RpcClient
        c = RpcClient(os.environ['IDATUI_RPC_SOCK'], timeout=60)
        c.connect()
        ok = 0
        for ea in sorted(cands):
            try:
                c.call('define', kind='func', target=hex(ea))
                ok += 1
            except Exception as e:
                print('   define %#x failed: %s' % (ea, str(e)[:50]))
        c.call('save')
        print('   defined %d/%d' % (ok, len(cands)))
    return 0


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