#!/usr/bin/env python3
"""一次性抓取: App 存储的 cred + Common/QUrl 两种 wupData 完整输出。"""
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 = r"""
'use strict';

function b64ify(ptr, n) {
  var bytes = ptr.readByteArray(n);
  var arr = new Uint8Array(bytes);
  var B64 = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
  var chars = "";
  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;
}

// otp 加密链 hook(二进制安全)
var base = Process.getModuleByName("libudbauthunify.so").base;
Interceptor.attach(base.add(0x32fa24), {
  onEnter: function (args) {
    function rd(p) {
      try {
        var b0 = p.readU8();
        if ((b0 & 1) === 0) return {t:"s", v:p.add(1).readUtf8String(b0>>1)};
        var len = parseInt(p.add(8).readU64().toString());
        return {t:"b64", v:b64ify(p.add(16).readPointer(), Math.min(len,4096))};
      } catch(e) { return {t:"err", v:String(e)}; }
    }
    this.s1 = rd(args[0]); this.b2 = args[2].toInt32();
    this.s2 = rd(args[3]); this.s3 = rd(args[4]); this.s4 = rd(args[5]);
    this.nonce = args[7].toString();
    this.outPtr = this.context.sp.readPointer();
  },
  onLeave: function () {
    function rdo(p) {
      try {
        var b0 = p.readU8();
        if ((b0 & 1) === 0) return {t:"b64", v:b64ify(p.add(1), b0>>1)};
        var len = parseInt(p.add(8).readU64().toString());
        return {t:"b64", v:b64ify(p.add(16).readPointer(), Math.min(len,4096))};
      } catch(e){ return {t:"err", v:String(e)}; }
    }
    send({type:"otp", s1:this.s1, b2:this.b2, s2:this.s2, s3:this.s3,
          s4:this.s4, nonce:this.nonce, out:rdo(this.outPtr)});
  }
});

function trigger() {
  Java.perform(function () {
    var attempt = 0;
    var timer = null;
    function once() {
      attempt += 1;
      var done = false;
      Java.enumerateClassLoaders({
        onMatch: function (loader) {
          if (done) return;
          try {
            var f = Java.ClassFactory.get(loader);
            var LP = f.use('com.hysdkproxy.LoginProxy');
            var inst = LP.getInstance();
            var HAS = f.use('com.huyaudbunify.HuyaAccountSaveUtils').getInstance();
            var uid = HAS.getUid();
            send({type:'uid', uid: String(uid)});
            var HA = f.use('com.huyaudbunify.HuyaAuth').getInstance();
            var cred = HA.getCred(uid);
            if (cred !== null) {
              send({type:'cred',
                    hyCred: cred.getHyCred() ? cred.getHyCred() : null,
                    yyCred: cred.getYyCred() ? cred.getYyCred() : null});
            } else {
              send({type:'cred', hyCred:null, yyCred:null});
            }
            var common = inst.getH5InfoEx();
            send({type:'wup_common', len: common?common.length:0, data: common});
            var qurl = inst.getQUrlData(uid, "", "");
            send({type:'wup_qurl', len: qurl?qurl.length:0, data: qurl});
            done = true;
          } catch (e) { /* 下一个 loader */ }
        },
        onComplete: function () {
          if (!done) {
            if (attempt % 4 === 0) send({type:'retry', n: attempt});
            timer = setTimeout(once, 3000);
          } else {
            send({type:'all_done', attempts: attempt});
          }
        }
      });
    }
    once();
  });
}

trigger();
"""


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--duration", type=float, default=40.0)
    ap.add_argument("--stabilize", type=float, default=10.0)
    ap.add_argument("--out", default="/tmp/full_dump.json")
    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)
    pid = device.spawn([args.package])
    print(f"spawned pid={pid}")
    session = device.attach(pid)
    session.create_script(bypass_src).load()
    device.resume(pid)
    time.sleep(args.stabilize)
    session.create_script(patch_guard).load()
    print("guards loaded")

    collected = {}

    def on_message(message, _data):
        if message.get("type") != "send":
            if message.get("type") == "error":
                print("ERR:", str(message)[:200])
            return
        p = message.get("payload") or {}
        t = p.get("type")
        if t in ("wup_common", "wup_qurl"):
            collected[t] = p.get("data")
            print(f"{t}: len={p.get('len')}")
        elif t == "cred":
            collected["cred"] = p
            hc, yc = p.get("hyCred"), p.get("yyCred")
            print(f"cred: hyCred={str(hc)[:60]}... ({len(hc) if hc else 0}) "
                  f"yyCred={str(yc)[:60]}... ({len(yc) if yc else 0})")
        elif t == "uid":
            collected["uid"] = p.get("uid")
            print(f"uid={p.get('uid')}")
        elif t == "loader_err":
            print("loader:", p.get("err"))
        else:
            print(json.dumps(p, ensure_ascii=False)[:120])

    script = session.create_script(JS)
    script.on("message", on_message)
    script.load()
    time.sleep(args.duration - args.stabilize)

    out = Path(args.out)
    out.write_text(json.dumps(collected, ensure_ascii=False))
    print(f"\nsaved: {out} keys={list(collected.keys())}")
    return 0


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