#!/usr/bin/env python3
"""真机: dlopen 拦截所有目标库, JNI_OnLoad 入口挂 RegisterNatives 钩子,
抓 NativeBridge.klog/b 等全部 native 注册归属(库+偏移).

用法: python3 scripts/phone_probe_rn_all.py [serial] [remote]
"""
from __future__ import annotations
import json, subprocess, sys, time
from pathlib import Path
import frida

SERIAL = sys.argv[1] if len(sys.argv) > 1 else "5dd8c93f"
REMOTE = sys.argv[2] if len(sys.argv) > 2 else "127.0.0.1:31878"
PACKAGE = "com.duowan.kiwi"
RE = Path("/Users/yml/codes/Reverse-Engineering-Agent-Universal-v3.0")
ART_CALLSITE = RE / "evidence" / "scripts" / "bypass_msaoaid_maps_art_callsite.js"

JS = r"""
'use strict';
function jmod(a){
  try{ var m = Process.findModuleByAddress(a); return m ? {n:m.name, off:'0x'+a.sub(m.base).toString(16)} : {n:'?',off:'?'}; }catch(e){ return {n:'?',off:'?'}; }
}
var rnHooked = false;
function hookRN(){
  if(rnHooked) return;
  try{
    var env = Java.vm.tryGetEnv();
    if(!env) return;
    var fns = env.handle.readPointer();
    var rn = fns.add(215 * Process.pointerSize).readPointer();
    if(rn.isNull()) return;
    Interceptor.attach(rn, {
      onEnter: function(args){
        var methods = args[3];
        var count = args[4].toInt32();
        var mm=[];
        for(var i=0;i<count;i++){
          var m = methods.add(i * Process.pointerSize * 3);
          try{
            var name = m.readPointer().readCString();
            var sig = m.add(Process.pointerSize).readPointer().readCString();
            var fnp = m.add(Process.pointerSize*2).readPointer();
            var jm = jmod(fnp);
            mm.push({name:name, sig:sig, mod:jm.n, off:jm.off});
          }catch(e){}
        }
        send({type:'rn', methods:mm, tid:Process.getCurrentThreadId()});
      }
    });
    rnHooked = true;
    send({type:'rn-hooked'});
  }catch(e){ send({type:'err', e:String(e).slice(0,120)}); }
}
setInterval(hookRN, 100);

// 目标库加载即挂 JNI_OnLoad
var TARGETS = ['libhydeviceid', 'libudbauthunify', 'libudb'];
var loaded = {};
function hookLib(md){
  if(loaded[''+md.base]) return;
  loaded[''+md.base] = true;
  var jni = md.findExportByName('JNI_OnLoad');
  if(!jni) return;
  var name = md.name;
  send({type:'lib', name:name, base:''+md.base});
  Interceptor.attach(jni, {
    onEnter: function(){ send({type:'ol-enter', lib:name}); hookRN(); },
    onLeave: function(ret){ send({type:'ol-leave', lib:name, ret:ret.toString()}); }
  });
}
setInterval(function(){
  Process.enumerateModules().forEach(function(md){
    for(var i=0;i<TARGETS.length;i++){
      if(md.name.indexOf(TARGETS[i]) >= 0){ hookLib(md); break; }
    }
  });
}, 100);
"""


def adb(*a):
    return subprocess.run(["adb", "-s", SERIAL, *a], capture_output=True, text=True)


def main():
    adb("shell", "am", "force-stop", PACKAGE)
    time.sleep(1.5)
    d = frida.get_device_manager().add_remote_device(REMOTE)
    pid = d.spawn([PACKAGE])
    print(f"[*] spawn {pid}", flush=True)
    s = d.attach(pid)
    s.create_script(ART_CALLSITE.read_text()).load()
    seen = set()
    regs = []

    def on_m(m, data):
        p = m.get("payload") or {}
        t = p.get("type")
        if t == "rn":
            for mm in p.get("methods", []):
                k = (mm["name"], mm["sig"], mm["mod"], mm["off"])
                if k not in seen:
                    seen.add(k)
                    regs.append(mm)
                    print(f"[RN] {mm['name']} {mm['sig']} @ {mm['mod']} {mm['off']}", flush=True)
        elif t in ("rn-hooked",):
            print("[*] rn-hooked", flush=True)
        elif t == "lib":
            print(f"[*] lib {p['name']} base={p['base']}", flush=True)
        elif t in ("ol-enter", "ol-leave"):
            print(f"[*] {t} {p.get('lib')} {p.get('ret','')}", flush=True)
        elif t == "err":
            print("[!]", p, flush=True)
        elif m.get("type") == "error":
            print("[JS-ERR]", m.get("description")[:120], flush=True)

    sc = s.create_script(JS)
    sc.on("message", on_m)
    sc.load()
    d.resume(pid)
    time.sleep(20)
    print(f"[*] total unique: {len(regs)}")


if __name__ == "__main__":
    main()