#!/usr/bin/env python3
"""nonce溯源v2: 修正random_device返回值读取 + libc熵源 + 同进程双调用对比。"""
from __future__ import annotations

import json
import subprocess
import time
from pathlib import Path

import frida

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

JS = r"""
'use strict';
Process.setExceptionHandler(function(d){
  send({type:'segv', info:{type:d.type, addr:String(d.address)}}); return true;
});
var base = Process.getModuleByName('libudbauthunify.so').base;
send({type:'armed', base:String(base)});
var WIN_ON = false;

// ---- random_device::operator() 返回uint32 ----
try {
  var sym = null;
  Process.getModuleByName('libudbauthunify.so').enumerateSymbols().forEach(function(s){
    if (s.name === '_ZNSt6__ndk113random_deviceclEv') sym = s.address;
  });
  if (sym) Interceptor.attach(sym, {
    onLeave: function(r){
      if (!WIN_ON) return;
      try { send({type:'rd32', v:'0x'+r.toInt32().toString(16), bt:DebugSymbol.fromAddress(this.returnAddress).toString().slice(0,60)}); } catch(e){}
    }
  });
} catch(e){ send({type:'hookerr', where:'rd_op', msg:''+e}); }

// ---- libc 熵源 ----
['getrandom','getentropy','arc4random_buf'].forEach(function(fn){
  var p = Module.findExportByName('libc.so', fn);
  if (!p) return;
  Interceptor.attach(p, {
    onEnter: function(a){ this.buf=a[0]||a[1]; this.n=a[1] ? a[1].toInt32() : a[0].toInt32(); },
    onLeave: function(_){
      if (!WIN_ON || this.n<=0 || this.n>64) return;
      try {
        var b = Memory.readByteArray(this.buf, Math.min(this.n,32));
        send({type:'libc', fn:fn, n:this.n,
              hex:Array.from(new Uint8Array(b)).map(function(x){return ('0'+x.toString(16)).slice(-2)}).join('')});
      } catch(e){}
    }
  });
});
var rd = Module.findExportByName('libc.so','read');
Interceptor.attach(rd, {onEnter:function(a){ this.fd=a[0].toInt32(); this.buf=a[1]; this.n=a[2].toInt32(); },
  onLeave:function(_){ if (!WIN_ON) return; }});

// ---- 双调用触发 ----
Java.perform(function(){
  var n=0;
  function go(){
    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();
      WIN_ON = true;
      send({type:'stage', msg:'call #1'});
      var q1=inst.getQUrlData(__UID__, "", "");
      send({type:'qurl_done', idx:1, data:q1||''});
      var q2=inst.getQUrlData(__UID__, "", "");
      WIN_ON = false;
      send({type:'qurl_done', idx:2, data:q2||''});
    }catch(e){ n+=1; WIN_ON=false; if(n%4===0) send({type:'retry',n:n,err:String(e).slice(0,90)}); setTimeout(go,5000); }
  }
  setTimeout(go, 20000);
});
""".replace('__UID__', str(UID))


def main():
    d = frida.get_device_manager().add_remote_device('127.0.0.1:31877')
    pid = d.spawn(['com.duowan.kiwi'])
    print(f'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/mask_frida_maps_only.js").read_text()).load()
    d.resume(pid)
    print('resumed')
    time.sleep(11)
    s.create_script((RE / "evidence/scripts/patch_guard_block_termination.js").read_text()).load()
    print('patch_guard on')

    sc = s.create_script(JS)
    events = []
    got_armed = []

    def on_msg(m, _):
        if m.get('type') == 'error':
            print('JS ERR:', str(m)[:200]); return
        if m.get('type') != 'send':
            return
        p = m.get('payload') or {}
        t = p.get('type')
        events.append(p)
        if t == 'armed':
            got_armed.append(1); print('[armed]')
        elif t == 'rd32':
            print(f"[rd32] {p['v']} <- {p['bt'][:50]}")
        elif t == 'libc':
            print(f"[{p['fn']}] n={p['n']} {p['hex'][:40]}")
        elif t == 'stage':
            print('[stage]', p['msg'])
        elif t == 'qurl_done':
            print(f"[qurl#{p.get('idx')}] len={p.get('len')} head={(p.get('data') or '')[:40]}")
        elif t == 'retry':
            print('[retry]', p.get('n'))

    sc.on('message', on_msg)
    sc.load()
    assert got_armed

    deadline = time.time() + 100
    while time.time() < deadline:
        time.sleep(2)
        if sum(1 for e in events if e.get('type') == 'qurl_done') >= 2:
            time.sleep(5)
            break
    Path('/tmp/opencode/nonce_v2.json').write_text(json.dumps(events))
    from collections import Counter
    print('saved:', Counter(e.get('type') for e in events))


if __name__ == '__main__':
    for attempt in range(3):
        print(f'===== 尝试 #{attempt+1} =====')
        try:
            main()
            break
        except frida.InvalidOperationError as e:
            print(f'session detached: {e}')
            subprocess.run(['adb', 'shell', 'am', 'force-stop', 'com.duowan.kiwi'])
            time.sleep(3)
        except Exception as e:
            print(f'err: {e}')
            subprocess.run(['adb', 'shell', 'am', 'force-stop', 'com.duowan.kiwi'])
            time.sleep(3)