#!/usr/bin/env python3
"""扫描全部 classes*.dex 的 code, 找调用指定方法的调用点 (调用者类/方法/code_off)。"""
import struct
import sys
from pathlib import Path
import dex_probe_native as d  # 复用 Dex 解析

INVOKE35 = {0x6e, 0x6f, 0x70, 0x71, 0x72, 0x74, 0x75, 0x76, 0x77, 0x78,
            0xf0, 0xf1, 0xf2, 0xf4, 0xf5, 0xf6, 0xf7}
INVOKE3RC = {0x6b, 0x6c, 0x6d, 0xfa, 0xfb, 0xfc, 0xfd}

def scan_code(dex, code_off, target_idx):
    if code_off == 0:
        return []
    off = code_off
    regs, ins_, outs, tries, dbg, insns = struct.unpack_from('<HHHHII', dex.data, off)
    insns_off = off + 16
    b = dex.data
    i = 0
    found = []
    while i < insns * 2:
        if insns_off + i + 2 > len(b):
            break
        op = b[insns_off + i]
        if op in INVOKE35:
            method_idx = struct.unpack_from('<H', b, insns_off + i + 2)[0]
            if method_idx == target_idx:
                found.append(insns_off + i)
            i += 4
        elif op in INVOKE3RC:
            method_idx = struct.unpack_from('<H', b, insns_off + i + 2)[0]
            if method_idx == target_idx:
                found.append(insns_off + i)
            i += 6
        else:
            # 近似长度: 30/31->1, 11x->1, 12x->1, 21x->2, 22x->2, 23x->2, 22c->2, 21c->2, 23c->2,
            # 20t->2, 22t->2, 21t->2, 22s->2, 21s->2, 2ss? -> 简单用 opcode 高低位判断
            hi = op >> 4
            lo = op & 0xf
            if hi == 0x1:  # 12x,11x,10x,0x
                i += 2
            elif hi in (0x2,):  # 21x,22x,23x,22c,21c,23c,22b,21b,22s,21s,21t,22t,20t,2(hi=2)
                i += 4
            elif hi == 0x3:  # 30t,32x,31i,31t,31c 或 3rc(op=0x3a-0x3d 但已处理部分)
                i += 6
            elif op in (0x0, 0xe):
                i += 2
            else:
                # 保守: 2 units
                i += 2
    return found

def main():
    dex_dir = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("work/apk_extract")
    target_cls = sys.argv[2] if len(sys.argv) > 2 else "Lcom/hy/HyDeviceProxy;"
    target_name = sys.argv[3] if len(sys.argv) > 3 else "j"
    target_proto = sys.argv[4] if len(sys.argv) > 4 else "V"  # 简化: 用 return type 匹配

    for dex_path in sorted(dex_dir.glob("classes*.dex")):
        data = dex_path.read_bytes()
        dex = d.Dex(data)
        # 找目标 method_idx
        t_idx = None
        for i in range(dex.method_ids_size):
            mcls, mname, pidx = dex.method(i)
            if mcls == target_cls and mname == target_name:
                rt, params = dex.proto(pidx)
                if "".join(params) + rt == target_proto or target_proto == "*":
                    t_idx = i
                    print(f"[{dex_path.name}] target method_idx={i} {mcls}.{mname}{tuple(params)}{rt}")
        if t_idx is None:
            continue
        for cd in dex.class_defs_iter():
            try:
                cls = dex.type(cd["class_idx"])
                methods = dex.class_data_methods(cd["class_data_off"])
            except Exception:
                continue
            for m_idx, acc, code_off in methods:
                if m_idx >= dex.method_ids_size or code_off == 0:
                    continue
                mcls, mname, pidx = dex.method(m_idx)
                f = scan_code(dex, code_off, t_idx)
                if f:
                    rt, params = dex.proto(pidx)
                    print(f"  CALLER {cls}.{mname}{tuple(params)}{rt} code_off={code_off:#x} calls@{[hex(x) for x in f]}")

if __name__ == "__main__":
    main()