#!/usr/bin/env python3
"""抓取证书加密真实密钥: hook UdbAESUtil::C1(构造器,收16B key) + encrypt 出入口。

背景: 差分已证明 Cipher=标准AES-128-ECB(tools/udb_aes.py selftest)，
剩下唯一未知 = 生产者每次加密用的 key 字节。C1 是所有路径的必经点。

时序(不可变): force-stop -> spawn挂起 -> 双bypass -> resume -> 10s -> patch_guard -> hooks
触发: Java桥直呼 LoginProxy.getQUrlData(uid,"","") 复现证书生成。
"""
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)});

function rdStr(p){ // libc++ string 安全读取, 返回 {t,len,v(hex)}
  try{
    var b0=p.readU8();
    if((b0&1)===0){ var l=b0>>1; if(l>4096) return {t:'s?',len:l,v:''};
      return {t:'s',len:l,v:l?Array.from(new Uint8Array(p.add(1).readByteArray(l))).map(x=>('0'+x.toString(16)).slice(-2)).join(''):''}; }
    var len=parseInt(p.add(8).readU64().toString());
    if(len>65536) return {t:'l?',len:len,v:''};
    var dp=p.add(16).readPointer();
    return {t:'l',len:len,v:Array.from(new Uint8Array(dp.readByteArray(len))).map(x=>('0'+x.toString(16)).slice(-2)).join('')};
  }catch(e){ return {t:'err',len:-1,v:String(e)}; }
}
function hx(p,n){ try{ return Array.from(new Uint8Array(p.readByteArray(n))).map(x=>('0'+x.toString(16)).slice(-2)).join(''); }catch(e){ return 'ERR'; } }

// 1) 构造器: 每次AES对象创建的16字节真钥
Interceptor.attach(base.add(0x24eb08), {
  onEnter: function(a){ send({type:'c1', key:hx(a[1],16), bt:DebugSymbol.fromAddress(this.returnAddress).toString().slice(0,80)}); }
});

// 2) 公开 encrypt(ret,x1,x2,...): 全参原始dump, 不再猜测约定
Interceptor.attach(base.add(0x250038), {
  onEnter: function(a){
    this.ra = DebugSymbol.fromAddress(this.returnAddress).toString().slice(0,90);
    this.retbuf = a[0];   // sret 缓冲指针(onLeave时 x0 可能已变, 用入口快照)
    send({type:'enc_in', ra:this.ra,
          x0s:rdStr(a[0]), x1s:rdStr(a[1]), x2s:rdStr(a[2]), x3s:rdStr(a[3]),
          x1raw:hx(a[1],32), x2raw:hx(a[2],32), x3raw:hx(a[3],32)});
  },
  onLeave: function(r){ send({type:'enc_out', ret:rdStr(this.retbuf)}); }
});

// 2.5) encode_aes(in,key,out): 看选key逻辑
Interceptor.attach(base.add(0x330218), {
  onEnter: function(a){
    this.p0=rdStr(a[0]); this.p1=rdStr(a[1]);
    this.ra = DebugSymbol.fromAddress(this.returnAddress).toString().slice(0,90);
  },
  onLeave: function(){ send({type:'encode_aes', ra:this.ra, p0:this.p0, p1:this.p1}); }
});

// 3) _encrypt(block,out&) 入口: 确认其调用方与块内容
Interceptor.attach(base.add(0x24f9f0), {
  onEnter: function(a){
    send({type:'_enc', ra:DebugSymbol.fromAddress(this.returnAddress).toString().slice(0,90),
          blk:hx(a[1],16)});
  }
});

// 自动触发证书生成
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();
      var q=inst.getQUrlData(__UID__, "", "");
      send({type:'qurl_done', len:q?q.length:0, data:q||''});
    }catch(e){ n+=1; if(n%6===0) send({type:'retry',n:n}); 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)[:240]); return
        if m.get('type') != 'send':
            return
        p = m.get('payload') or {}
        t = p.get('type')
        if t != 'retry':
            events.append(p)
        if t == 'armed':
            got_armed.append(1); print('[armed]', p['base'])
        elif t == 'c1':
            print(f"[C1] key={p['key']}  <- {p['bt'][:60]}")
        elif t == 'enc_in':
            print(f"[enc] ra={p['ra'][:60]}")
            print(f"   x1s={p['x1s']['t']}:{p['x1s']['len']} head={p['x1s']['v'][:48]}")
            print(f"   x2s={p['x2s']['t']}:{p['x2s']['len']} head={p['x2s']['v'][:48]}")
            print(f"   x3s={p['x3s']['t']}:{p['x3s']['len']} head={p['x3s']['v'][:48]}")
        elif t == 'enc_out':
            print(f"   out={p['ret']['t']}:{p['ret']['len']} head={p['ret']['v'][:48]}")
        elif t == '_enc':
            print(f"[_encrypt] blk={p['blk']} <- {p['ra'][:60]}")
        elif t == 'encode_aes':
            print(f"[encode_aes] p0={p['p0']['t']}:{p['p0']['len']},{p['p0']['v'][:48]} "
                  f"p1={p['p1']['t']}:{p['p1']['len']},{p['p1']['v'][:48]} <- {p['ra'][:56]}")
        elif t == 'qurl_done':
            print('[qurl_done]', p.get('len'), 'head:', (p.get('data') or '')[:60])

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

    deadline = time.time() + 150
    while time.time() < deadline:
        time.sleep(2)
        if any(e.get('type') == 'qurl_done' for e in events):
            time.sleep(8)
            break

    Path('/Users/yml/codes/douyu_login_py/evidence/cert_keycap.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(4):
        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}')
            time.sleep(3)