#!/usr/bin/env python3
"""UdbAESUtil 差分分析 harness（NativeFunction 直调）。

时序（不可变）: force-stop -> spawn挂起 -> bypass -> resume -> 稳定10s -> patch_guard -> 业务脚本
静态已确认对象布局: +0x00 vptr | +0x08 S-box拷贝(256B) | +0x108 InvS-box拷贝(256B) | +0x208 轮密钥(176B)

导出符号（.dynsym）:
  C1(uchar*key)              0x24eb08   构造器
  KeyExpansion               0x24ebd4
  Cipher(uchar*)             0x24f314   10轮
  InvCipher(uchar*)          0x24f778
  _encrypt(uchar*,string&)   0x24f9f0
  _decrypt                   0x24fe90
  encrypt(ret&,key?,in?)     0x250038
  decrypt                    0x2501a4
"""
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")

JS = r"""
'use strict';
// 段错误防护: 拦截 native 崩溃, 上报后恢复而不是带崩进程
Process.setExceptionHandler(function(det){
  send({type:'segv', info:{type:det.type, addr:String(det.address), msg:String(det.memory||'')}});
  return true;
});
var MOD = 'libudbauthunify.so';
var base = Process.getModuleByName(MOD).base;
send({type:'armed', base:String(base)});

var A = {
  C1:        new NativeFunction(base.add(0x24eb08), 'void', ['pointer','pointer']),
  Cipher:    new NativeFunction(base.add(0x24f314), 'void', ['pointer','pointer']),
  InvCipher: new NativeFunction(base.add(0x24f778), 'void', ['pointer','pointer']),
  enc_pub:   new NativeFunction(base.add(0x250038), 'pointer', ['pointer','pointer','pointer']),
  dec_pub:   new NativeFunction(base.add(0x2501a4), 'pointer', ['pointer','pointer','pointer']),
};

function hex(p,n){ return Array.from(new Uint8Array(p.readByteArray(n))).map(x=>('0'+x.toString(16)).slice(-2)).join(''); }
function unhex(s){ var n=s.length>>1, b=new Uint8Array(n); for(var i=0;i<n;i++) b[i]=parseInt(s.substr(i*2,2),16); return b; }
function wrBytes(p, arr){ p.writeByteArray(arr); }

// libc++ std::string 构造 (老布局: cap@0|flagLSB, size@8, data@16 / SSO: b0=len<<1,data@1)
function mkstr(bytes){
  var m = Memory.alloc(32);
  var n = bytes.length;
  if (n <= 22) {
    Memory.writeByteArray(m, new Array(24).fill(0));
    m.writeU8(n<<1);
    if(n) wrBytes(m.add(1), bytes);
  } else {
    var buf = Memory.alloc(n+1);
    wrBytes(buf, bytes); buf.add(n).writeU8(0);
    m.writeU64( uint64(((n)<<1)|1) );
    m.add(8).writeU64(uint64(n));
    m.add(16).writePointer(buf);
  }
  return m;
}
function strBytes(sp){ // 读回 std::string 内容
  var b0=sp.readU8();
  if((b0&1)===0){ var l=b0>>1; return l? new Uint8Array(sp.add(1).readByteArray(l)) : new Uint8Array(0); }
  var len=parseInt(sp.add(8).readU64().toString());
  var dp=sp.add(16).readPointer();
  return len? new Uint8Array(dp.readByteArray(len)) : new Uint8Array(0);
}

var OBJLEN = 0x300;
function newObj(){ return Memory.alloc(OBJLEN); }

rpc.exports = {
  // 轮密钥展开对拍
  sched: function(keyHex){
    var o=newObj(), k=Memory.alloc(32);
    wrBytes(k, unhex(keyHex));
    A.C1(o, k);
    return { rk: hex(o.add(0x208), 176), sbox_head: hex(o.add(8),16) };
  },
  // 单块 Cipher（读改前后，判断是否就地）
  cipher: function(keyHex, blkHex){
    var o=newObj(), k=Memory.alloc(32), b=Memory.alloc(48);
    wrBytes(k, unhex(keyHex)); wrBytes(b, unhex(blkHex));
    A.C1(o,k);
    var before = hex(b,16);
    A.Cipher(o,b);
    return { before:before, after:hex(b,16) };
  },
  invcipher: function(keyHex, ctHex){
    var o=newObj(), k=Memory.alloc(32), b=Memory.alloc(48);
    wrBytes(k, unhex(keyHex)); wrBytes(b, unhex(ctHex));
    A.C1(o,k);
    A.InvCipher(o,b);
    return { after: hex(b,16) };
  },
  // 公开 encrypt(ret,this=x1,x2) —— 按静态结论: static(ret=&ret, x1=&plain, x2=key临时)
  encpub: function(plainHex, keyBytesHex){
    var ret=Memory.alloc(32), plain=mkstr(unhex(plainHex)), key=mkstr(unhex(keyBytesHex));
    try {
      A.enc_pub(ret, plain, key);
      return { out: hex(strBytes(ret)) };
    } catch(e){ return { err:String(e) }; }
  },
  decpub: function(ctHex, keyBytesHex){
    var ret=Memory.alloc(32), ct=mkstr(unhex(ctHex)), key=mkstr(unhex(keyBytesHex));
    try {
      A.dec_pub(ret, ct, key);
      return { out: hex(strBytes(ret)) };
    } catch(e){ return { err:String(e) }; }
  },
  ping: function(){ return 'pong'; }
};
"""


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)
    # 双重防护(与 hook_final_capture 完全一致): callsite bypass + maps 掩盖
    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)
    state = {'armed': False, 'segv': 0}
    OUT = Path('/Users/yml/codes/douyu_login_py/evidence/aes_diff.json')
    def save(k, v):
        try:
            data = json.loads(OUT.read_text()) if OUT.exists() else {}
        except Exception:
            data = {}
        data[k] = v
        OUT.write_text(json.dumps(data, indent=1))
    def on_msg(m, _):
        if m.get('type') == 'error':
            print('JS ERR:', str(m)[:300])
        elif m.get('type') == 'send':
            p = m.get('payload') or {}
            if p.get('type') == 'armed':
                state['armed'] = True
                print('[armed]', p['base'])
            elif p.get('type') == 'segv':
                state['segv'] += 1
                print(f"[SEGV #{state['segv']}] {p['info']}")
    sc.on('message', on_msg)
    sc.load()
    assert state['armed']
    print('ping:', sc.exports.ping())

    k16 = '000102030405060708090a0b0c0d0e0f'
    KEY27 = 'HuyaUdb1928374650qwertyuiop'.encode().hex()

    # ---- 阶段1: KeyExpansion 轮密钥对拍 ----
    r = sc.exports.sched(k16)
    save('sched_kat', r); print('[sched] rk[:32]', r['rk'][:32])

    # ---- 阶段2: 单块 Cipher KAT ----
    r = sc.exports.cipher(k16, '00112233445566778899aabbccddeeff')
    save('cipher_kat', r)
    std = '69c4e0d86a7b0430d8cdb78070b4c55a'
    ok = r.get('after') == std
    print(f'[cipher KAT] {r.get("after")} 标准FIPS-197={std} match={ok}')
    if not ok:
        print('!! 与标准不一致 -> 自定义点在轮结构/密钥扩展, 后续阶段仍执行')

    # ---- 阶段3: HuyaUdb 主密钥 first16 差分 ----
    r = sc.exports.cipher(KEY27.encode().hex()[:32], '00' * 16)
    save('cipher_huyaudb_first16', r); print('[cipher first16(HuyaUdb)]', r)

    # ---- 阶段4: InvCipher 往返 ----
    ct = sc.exports.cipher(k16, 'cafebabe' * 4).get('after', '')
    pt = sc.exports.invcipher(k16, ct).get('after', '') if ct else ''
    print(f'[inv roundtrip] ct={ct[:32]} -> pt={pt} match={pt == "cafebabe" * 4}')
    save('inv_roundtrip', {'ct': ct, 'pt': pt})

    # ---- 阶段5: 公开 encrypt(最可能崩, 放最后) ----
    for name, pt in [('empty', ''), ('15B', 'aa' * 15), ('16B', 'bb' * 16), ('187B', 'cc' * 187)]:
        try:
            r = sc.exports.encpub(pt, KEY27)
            o = r.get('out', '')
            print(f'[encpub {name}] -> {len(o)//2}B head={o[:48]}')
            save(f'encpub_{name}', r)
        except Exception as e:
            print(f'[encpub {name}] EXC {e}')

    try:
        e = sc.exports.encpub('ab' * 40, KEY27)
        dd = sc.exports.decpub(e.get('out', ''), KEY27)
        print(f'[roundtrip 80B] ct={len(e.get("out",""))//2}B pt_ok={dd.get("out","").startswith("ab"*16)}')
        save('roundtrip80', {'e': e, 'd': dd})
    except Exception as ex:
        print('[roundtrip] EXC', ex)

    print('done, segv count:', state['segv'])


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)