#!/usr/bin/env python3
"""dump 魔数缓冲区: 找到魔数地址后, dump 前后完整内存到证据文件。

命中即写入文件(不依赖退出), 打印明文头与密文长度。
"""
from __future__ import annotations

import time
from pathlib import Path

import frida

RE = Path("/Users/yml/codes/Reverse-Engineering-Agent-Universal-v3.0")
OUT = Path("/Users/yml/codes/douyu_login_py/evidence/dfp_buffer_dump.json")

JS = r"""
'use strict';
Process.setExceptionHandler(function(d){ return true; });
send({type:'armed'});

function hex(p, n){
  try{ return Array.from(new Uint8Array(p.readByteArray(n))).map(x=>('0'+x.toString(16)).slice(-2)).join(''); }
  catch(e){ return 'ERR'; }
}

function doScan(){
  var found = false;
  try{
    Process.enumerateRanges('r--').forEach(function(r){
      if (r.size > 1024*1024*256) return;
      try{
        var m = Memory.scanSync(r.base, r.size, '57 18 82 cf 66 4b b3 94');
        m.forEach(function(x){
          var addr = x.address;
          // 只 dump 每个 4KB 页内第一个命中 (去重)
          var pre256 = hex(addr.sub(64), 64);
          var post = hex(addr.add(10), 4096);
          send({type:'dump', addr:String(addr), pre:pre256, post:post});
          found = true;
        });
      }catch(e){}
    });
  }catch(e){}
  send({type:'tick', found:found});
}

var iv = setInterval(function(){ doScan(); }, 3000);
"""


def main():
    device = frida.get_device_manager().add_remote_device("127.0.0.1:31877")
    pid = device.spawn(["com.duowan.kiwi"])
    print(f"[*] spawned pid={pid}", flush=True)
    session = device.attach(pid)
    session.create_script(
        (RE / "evidence/scripts/bypass_msaoaid_maps_art_callsite.js").read_text()
    ).load()
    session.create_script(
        (RE / "evidence/scripts/mask_frida_maps_only.js").read_text()
    ).load()
    device.resume(pid)
    time.sleep(11)
    session.create_script(
        (RE / "evidence/scripts/patch_guard_block_termination.js").read_text()
    ).load()

    dumped = []

    def on_message(message, data):
        if message.get("type") == "error":
            print("[JS-ERR]", str(message)[:160], flush=True); return
        if message.get("type") != "send":
            return
        p = message.get("payload") or {}
        t = p.get("type")
        if t == "dump":
            addr = p.get("addr")
            # 去重 (同地址不重复存)
            if any(d.get("addr") == addr for d in dumped):
                return
            rec = {"addr": addr, "pre": p.get("pre"), "post": p.get("post")}
            dumped.append(rec)
            # 立即持久化
            OUT.write_text(__import__("json").dumps(dumped, indent=1))
            print(f"[DUMP] @{addr} pre={p.get('pre','')[:64]} postlen={len(p.get('post',''))//2} saved={len(dumped)}", flush=True)
        elif t == "tick":
            print(f"[tick] found={p.get('found')} total_dumps={len(dumped)}", flush=True)

    script = session.create_script(JS)
    script.on("message", on_message)
    script.load()
    print("[*] dump 扫描中 (每3s), 命中即持久化。请触发登录 (命中后可 pkill)", flush=True)
    try:
        while True:
            time.sleep(5)
    except KeyboardInterrupt:
        pass
    print(f"[*] 结束, 共 {len(dumped)} 个 dump -> {OUT}", flush=True)


if __name__ == "__main__":
    main()