#!/usr/bin/env python3
"""最终轮 hook: 加密函数现场抓取(AES key/明文/密文) + Java 触发 getQUrlData。"""
from __future__ import annotations

import json
import time
from pathlib import Path

import frida

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

JS = r"""
'use strict';
function b64ify(ptr, n) {
  var arr = new Uint8Array(ptr.readByteArray(n));
  var B64 = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
  var c = "";
  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;
    c+=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 c;
}
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());
    if(len>8192) return {t:"err",v:"too long"};
    return {t:"b64",v:b64ify(p.add(16).readPointer(),len)};
  }catch(e){return {t:"err",v:String(e)};}
}

var armed = false;
var timer = setInterval(function(){
  if (armed) { clearInterval(timer); return; }
  var base;
  try { base = Process.getModuleByName("libudbauthunify.so").base; }
  catch(e) { return; }
  armed = true;
  clearInterval(timer);
  send({type:'base', v:String(base)});

  // md5_char16(string& src, string& out16)
  Interceptor.attach(base.add(0x32e71c), {
    onEnter: function(a){
      this.inp = rd(a[0]);
      this.outp = a[1];
      // 原始内存也存一份
      try { this.rawin = b64ify(a[0], 40); } catch(e){ this.rawin=''; }
    },
    onLeave: function(){
      send({type:'md5c16', inp:this.inp, out:rd(this.outp), rawin:this.rawin});
    }
  });

  // KeyExpansion: 两个参数都转储
  var keCount = 0;
  Interceptor.attach(base.add(0x24ebd4), {
    onEnter: function(a){
      keCount += 1;
      if (keCount > 2) return;
      try {
        var h = function(p,n){ return Array.from(new Uint8Array(p.readByteArray(n))).map(function(x){return ('0'+x.toString(16)).slice(-2);}).join(''); };
        send({type:'aeskey_raw', n:keCount, x0:h(a[0],32), x1:h(a[1],48)});
      } catch(e){}
    }
  });

  // UdbAESUtil::encrypt: x2 字符串对象的原始内存转储
  Interceptor.attach(base.add(0x250038), {
    onEnter: function(a){
      this.inp = rd(a[1]);
      try {
        var p = a[2];
        var b0 = p.readU8();
        var info = {b0:b0};
        if ((b0&1)===0) { info.len=b0>>1; info.data=b64ify(p.add(1), Math.min(b0>>1,48)); }
        else {
          info.len = parseInt(p.add(8).readU64().toString());
          info.cap = parseInt(p.readU64().toString());
          info.data = b64ify(p.add(16).readPointer(), Math.min(info.len,64));
        }
        send({type:'aeskey', info:info});
      } catch(e){ send({type:'aeskey', info:{err:String(e)}}); }
      this.out = a[3];
    },
    onLeave: function(){ send({type:'aes_out', out:rd(this.out)}); }
  });

  // hyudb_otp_encrypt
  Interceptor.attach(base.add(0x32fa24), {
    onEnter: function(a){
      this.s1=rd(a[0]); this.b1=a[1].toInt32(); this.b2=a[2].toInt32();
      this.s2=rd(a[3]); this.s3=rd(a[4]); this.s4=rd(a[5]);
      this.n=a[7].toString();
      this.outP=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.n,out:rd(this.outP)});
    }
  });
  send({type:'armed'});

  // 触发: 等类可用后轮询
  Java.perform(function(){
    var n = 0;
    function tryOnce(){
      try {
        var seed = Java.use('com.duowan.kiwi.base.login.udb.HuyaLoginProxy');
        var F = Java.ClassFactory.get(seed.class.getClassLoader());
        var inst = F.use('com.hysdkproxy.LoginProxy').getInstance();
        var q = inst.getQUrlData(1199666914671, "", "");
        send({type:'qurl_done', len: q?q.length:0});
        var c = inst.getH5InfoEx();
        send({type:'common_done', len: c?c.length:0});
      } catch(e) {
        n += 1;
        if (n % 6 === 0) send({type:'trig_retry', n:n, err:String(e).substring(0,70)});
        setTimeout(tryOnce, 5000);
      }
    }
    setTimeout(tryOnce, 3000);
  });
}, 400);
"""


def main():
    d = frida.get_device_manager().add_remote_device('127.0.0.1:31877')
    pid = d.spawn(['com.duowan.kiwi'])
    print('spawned', pid)
    s = d.attach(pid)
    s.create_script((RE / "evidence/scripts/bypass_msaoaid_maps_art_callsite.js").read_text()).load()
    s.create_script((RE / "evidence/scripts/patch_guard_block_termination.js").read_text()).load()
    d.resume(pid)
    print('resumed')

    events = []
    def on_msg(m, _):
        if m.get('type') != 'send':
            if m.get('type') == 'error':
                print('ERR:', str(m)[:180])
            return
        p = m.get('payload') or {}
        events.append(p)
        t = p.get('type')
        if t == 'otp':
            print(f"\nOTP s1={p['s1']['v'][:20]} b1={p['b1']} b2={p['b2']} "
                  f"s2={p['s2']['v']} s3={str(p['s3'])[:40]}")
            s4 = p['s4']
            print(f"   s4[{s4.get('t')}] {len(s4.get('v',''))}B nonce={p['nonce']}")

    sc = s.create_script(JS)
    sc.on('message', on_msg)
    sc.load()
    time.sleep(75)
    Path('/tmp/crypto_events.json').write_text(json.dumps(events, ensure_ascii=False))
    aes = [e for e in events if e.get('type') == 'aes_in']
    print(f"\n事件统计: otp={sum(1 for e in events if e.get('type')=='otp')} "
          f"aes={len(aes)} ctp={sum(1 for e in events if e.get('type')=='ctp')}")
    for e in aes[:6]:
        ki = e['key']
        kv = ki['v'] if isinstance(ki, dict) else str(ki)
        ii = e['inp']
        iv = ii['v'] if isinstance(ii, dict) else str(ii)
        print(f"AES key[{kv[:24]}...] in[{len(iv)}B]")


if __name__ == '__main__':
    main()