#!/usr/bin/env python3
"""追踪 this+0x10(k1) 的写入者: hook std::string::operator=, 比对目标地址==BusinessCfg+0x10。
启动后等初始化完成, 直接读 BusinessCfg 实例的 +0x10 最终值, 并记录所有写到该地址的调用。
"""
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';
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){
  try{
    var b0=p.readU8();
    if((b0&1)===0){ var l=b0>>1; if(l>8192) 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)}; }
}

// BusinessCfg::getInstance 返回单例 -> 拿 this 地址
var cfgAddr = null;
try {
  Interceptor.attach(base.add(0x454310 - 0x454310 + 0x26a880), {  // getInstance 附近, 用符号找
  });
} catch(e){}

// 用 export 查找 getInstance 真实地址
var syms = Process.getModuleByName('libudbauthunify.so').enumerateSymbols();
var getInst = null;
syms.forEach(function(s){ if (s.name.indexOf('BusinessCfg11getInstance') >= 0) getInst = s.address; });
send({type:'info', getInst:String(getInst)});

if (getInst) {
  Interceptor.attach(getInst, {
    onLeave: function(r){
      cfgAddr = r;
      var s = rdStr(r.add(0x10));
      send({type:'cfg', addr:String(r), k1: s});
    }
  });
}

// hook std::string::operator=(string const&) @0x452e40
try {
  Interceptor.attach(base.add(0x452e40), {
    onEnter: function(a){
      if (!cfgAddr) return;
      // a[0] = this(目标), a[1] = src
      if (a[0].equals(cfgAddr.add(0x10))) {
        var v = rdStr(a[1]);
        send({type:'write_k1', val:v,
              bt:Thread.backtrace(this.context, Backtracer.ACCURATE).map(function(x){
                return DebugSymbol.fromAddress(x).toString().slice(0,80);}).slice(0,4)});
      }
    }
  });
} catch(e){ send({type:'hookerr', where:'op=', msg:''+e}); }

// 直接读最终值
setTimeout(function(){
  if (cfgAddr) {
    var s = rdStr(cfgAddr.add(0x10));
    send({type:'final_k1', k1:s});
  }
}, 60000);
"""


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')
        events.append(p)
        if t == 'armed':
            got_armed.append(1); print('[armed]')
        elif t == 'info':
            print('[info] getInstance =', p['getInst'])
        elif t == 'cfg':
            print(f"[cfg] addr={p['addr']} k1={p['k1']['v'][:40]}({p['k1']['len']})")
        elif t == 'write_k1':
            print(f"[WRITE k1] val={p['val']['v'][:60]}")
            for f in p['bt']:
                print("    <", f)
        elif t == 'final_k1':
            print(f"[FINAL k1] {p['k1']['v'][:40]}({p['k1']['len']})")
        elif t == 'hookerr':
            print('[hookerr]', p)

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

    deadline = time.time() + 130
    while time.time() < deadline:
        time.sleep(2)
        if any(e.get('type') == 'final_k1' for e in events):
            time.sleep(3)
            break
    Path('/tmp/opencode/k1_writer.json').write_text(json.dumps(events))
    print('saved /tmp/opencode/k1_writer.json')


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)