#!/usr/bin/env python3
"""全自动 hook：触发并抓取虎牙 App 的 H5InfoEx(证书)生成全链路。

流程: spawn 挂起 → 构造窗口挂 bypass → resume → 稳定后挂业务 hooks
     → 通过 Java 桥直呼 LoginProxy.getH5InfoEx() 触发生成（零人工）
输出: /tmp/h5infoex_dump.jsonl
"""
from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import frida

RE_DIR = Path("/Users/yml/codes/Reverse-Engineering-Agent-Universal-v3.0")

JS_AGENT = r"""
'use strict';

function rdStr(p) {
  try {
    var b0 = p.readU8();
    if ((b0 & 1) === 0) {
      var n = b0 >> 1;
      return {t:"s", v: p.add(1).readUtf8String(n)};
    }
    var len = parseInt(p.add(8).readU64().toString());
    var dp = p.add(16).readPointer();
    return {t:"s", v: dp.readUtf8String(len)};
  } catch (e) {
    // 二进制内容: 走 base64
    try {
      var b0b = p.readU8();
      if ((b0b & 1) === 0) {
        var nb = b0b >> 1;
        return {t:"b64", v: base64ify(p.add(1), nb)};
      }
      var len2 = parseInt(p.add(8).readU64().toString());
      var dp2 = p.add(16).readPointer();
      return {t:"b64", v: base64ify(dp2, Math.min(len2, 4096))};
    } catch (e2) {
      return {t:"err", v: String(e)};
    }
  }
}

function base64ify(ptr, n) {
  var bytes = ptr.readByteArray(n);
  var arr = new Uint8Array(bytes);
  var chars = "";
  var B64 = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
  for (var i = 0; i < arr.length; i += 3) {
    var b0 = arr[i], b1 = i+1 < arr.length ? arr[i+1] : 0, b2 = i+2 < arr.length ? arr[i+2] : 0;
    chars += B64[b0>>2] + B64[((b0&3)<<4)|(b1>>4)] +
             (i+1 < arr.length ? B64[((b1&15)<<2)|(b2>>6)] : "=") +
             (i+2 < arr.length ? B64[b2&63] : "=");
  }
  return chars;
}

var base = Process.getModuleByName("libudbauthunify.so").base;
send({type:"info", msg:"base=" + base});

// hyudb_otp_encrypt(string,u8,u8,string,string,string,u8,ulong,string&)
Interceptor.attach(base.add(0x32fa24), {
  onEnter: function (args) {
    this.s1 = rdStr(args[0]);
    this.b1 = args[1].toInt32();
    this.b2 = args[2].toInt32();
    this.s2 = rdStr(args[3]);
    this.s3 = rdStr(args[4]);
    this.s4 = rdStr(args[5]);
    this.b3 = args[6].toInt32();
    this.nonce = args[7].toString();
    this.outPtr = this.context.sp.readPointer();
  },
  onLeave: function () {
    send({type:"otp", s1:this.s1, b1:this.b1, b2:this.b2, s2:this.s2,
          s3:this.s3, s4:this.s4, b3:this.b3, nonce:this.nonce,
          out: rdStr(this.outPtr)});
  }
});

// BusinessCfg::getOtp(u64, string&, int&)
Interceptor.attach(base.add(0x2689f4), {
  onEnter: function (args) {
    this.uid = args[1].toString();
    this.outStr = args[2];
  },
  onLeave: function () {
    send({type:"getOtp", uid:this.uid, out: rdStr(this.outStr)});
  }
});

// WupDataPackage<AppCommonData>::createWupRequestData(...)
Interceptor.attach(base.add(0x38dab0), {
  onEnter: function (args) {
    this.ret = args[0];
    this.c3 = args[3].readCString();
    this.c4 = args[4].readCString();
  },
  onLeave: function () {
    send({type:"createWup", servant:this.c3, func:this.c4,
          result_b64: rdStr(this.ret)});
  }
});

send({type:"ready"});
"""

JS_TRIGGER = r"""
'use strict';
function tryTrigger() {
  Java.perform(function () {
    var found = false;
    Java.enumerateClassLoaders({
      onMatch: function (loader) {
        if (found) return;
        try {
          var factory = Java.ClassFactory.get(loader);
          var LP = factory.use('com.hysdkproxy.LoginProxy');
          var inst = LP.getInstance();
          var r = inst.getH5InfoEx();
          found = true;
          send({type:'trigger_ok', len: r ? r.length : 0,
                head: r ? r.substring(0, 80) : 'null'});
        } catch (e) {
          send({type:'trigger_try_err', err: String(e).substring(0,120)});
        }
      },
      onComplete: function () {
        if (!found) send({type:'trigger_fail'});
      }
    });
  });
}
tryTrigger();
"""


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--duration", type=float, default=90.0)
    ap.add_argument("--stabilize", type=float, default=10.0)
    ap.add_argument("--out", default="/tmp/h5infoex_dump.jsonl")
    ap.add_argument("--package", default="com.duowan.kiwi")
    ap.add_argument("--port", default="127.0.0.1:31877")
    args = ap.parse_args()

    bypass_src = (RE_DIR / "evidence/scripts/bypass_msaoaid_maps_art_callsite.js").read_text()
    patch_guard = (RE_DIR / "evidence/scripts/patch_guard_block_termination.js").read_text()

    device = frida.get_device_manager().add_remote_device(args.port)
    print(f"remote device ok")

    # 1) spawn 挂起 → 构造窗口期只挂 bypass
    pid = device.spawn([args.package])
    print(f"spawned pid={pid}")
    session = device.attach(pid)
    s_bypass = session.create_script(bypass_src)
    s_bypass.load()
    device.resume(pid)
    print("bypass loaded, resumed")
    time.sleep(args.stabilize)

    # 2) patch_guard + 业务 hooks
    session.create_script(patch_guard).load()
    out_path = Path(args.out)
    fout = out_path.open("w")
    got = {"otp": False, "wup": False}

    def on_message(message, _data):
        if message.get("type") != "send":
            if message.get("type") == "error":
                print("SCRIPT ERR:", str(message)[:200])
            return
        p = message.get("payload") or {}
        fout.write(json.dumps(p, ensure_ascii=False) + "\n")
        fout.flush()
        t = p.get("type")
        def unwrap(x):
            if isinstance(x, dict) and "t" in x:
                v = x["v"]
                tag = x["t"]
                if isinstance(v, str) and len(v) > 200:
                    return f"[{tag}]{v[:150]}...({len(v)}B)"
                return f"[{tag}]{v}"
            return x
        if t == "otp":
            got["otp"] = True
            print("\n=== hyudb_otp_encrypt ===")
            for k in ("s1","b1","b2","s2","s3","s4","b3","nonce"):
                print(f"  {k} = {unwrap(p.get(k))}")
            print(f"  out = {unwrap(p.get('out'))}")
        elif t == "createWup":
            got["wup"] = True
            r = p["result_b64"]
            print(f"\n=== createWupRequestData servant={p['servant']} func={p['func']} ===")
            print(f"  result({len(r)}B) = {r[:140]}...")
        elif t == "getOtp":
            o = p.get("out")
            if isinstance(o, dict):
                o = f"[{o['t']}]{str(o['v'])[:140]}"
            print(f"\n=== getOtp uid={p['uid']} out={o}")
        elif t in ("ready", "info"):
            print(f"[{t}] {p.get('msg','')}")

    script = session.create_script(JS_AGENT)
    script.on("message", on_message)
    script.load()
    print("\nhooks armed")

    # 3) 轮询式 Java 触发（等 App 完全初始化）
    s_trig = session.create_script(JS_TRIGGER)
    s_trig.on("message", on_message)
    deadline = time.time() + args.duration
    triggered = False
    while time.time() < deadline and not triggered:
        try:
            s_t2 = session.create_script(JS_TRIGGER)
            s_t2.on("message", on_message)
            s_t2.load()
            time.sleep(6)
            if got["otp"] or got["wup"]:
                triggered = True
        except Exception as e:
            print("trigger retry:", str(e)[:80])
            time.sleep(3)

    # 等最后的异步 dump 落盘
    time.sleep(3)
    fout.close()
    print(f"\nresult: otp={got['otp']} wup={got['wup']} -> {out_path}")
    return 0 if (got["otp"] and got["wup"]) else 1


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