#!/usr/bin/env python3
"""
Frida 本地 attach 测试 — 验证 get_local_device() 捕获 SQLCipher key。
在 Ubuntu proot 容器内执行。
"""
import subprocess, sys, json, time, os

REPORT = {"test": "frida_local", "steps": {}, "overall": None}

def log_step(name, ok, detail=""):
    status = "[OK]" if ok else "[FAIL]"
    print(f"  {status} {name}")
    if detail:
        print(f"       {detail}")
    REPORT["steps"][name] = {"ok": ok, "detail": detail}
    return ok

def get_wechat_pid():
    """Get WeChat main process PID via pgrep."""
    try:
        r = subprocess.run(
            ["su", "-c", "pgrep -x com.tencent.mm"],
            capture_output=True, text=True, timeout=5
        )
        if r.returncode == 0 and r.stdout.strip():
            pids = sorted([int(p.strip()) for p in r.stdout.strip().split("\n") if p.strip()])
            return pids[0]  # Main process is smallest PID
    except Exception as e:
        print(f"       pgrep error: {e}")
    return None

def main():
    print("=" * 60)
    print("Frida 本地 Attach 测试")
    print("=" * 60)

    all_ok = True

    # Step 1: Check frida Python package
    print("\n[Step 1] Frida Python 包")
    try:
        import frida
        print(f"  [OK] frida {frida.__version__}")
        REPORT["steps"]["frida_python"] = {"ok": True, "version": frida.__version__}
    except ImportError:
        print("  [FAIL] frida 未安装。运行: pip install frida==17.9.10")
        REPORT["steps"]["frida_python"] = {"ok": False}
        save_report()
        return 1

    # Step 2: Check frida-server
    print("\n[Step 2] Frida Server 运行状态")
    r = subprocess.run(["su", "-c", "ps -A | grep frida-server"],
                       capture_output=True, text=True)
    has_frida_server = r.returncode == 0 and "frida-server" in r.stdout
    all_ok &= log_step("frida_server", has_frida_server,
                       "running" if has_frida_server else "frida-server 未运行")

    # Step 3: Get local device
    print("\n[Step 3] 获取本地 Frida 设备")
    try:
        import frida
        device = frida.get_local_device()
        log_step("local_device", True, f"device.name={device.name}")
    except Exception as e:
        all_ok &= log_step("local_device", False, str(e))
        save_report()
        return 1

    # Step 4: Find WeChat PID
    print("\n[Step 4] 查找微信进程 PID")
    pid = get_wechat_pid()
    if pid:
        all_ok &= log_step("wechat_pid", True, f"PID={pid}")
    else:
        print("       微信未运行，尝试启动...")
        subprocess.run(["su", "-c", "am start -n com.tencent.mm/.ui.LauncherUI"],
                       capture_output=True)
        time.sleep(3)
        pid = get_wechat_pid()
        if pid:
            all_ok &= log_step("wechat_pid", True, f"PID={pid} (已自动启动)")
        else:
            all_ok &= log_step("wechat_pid", False, "无法获取微信 PID")
            save_report()
            return 1

    # Step 5: Attach to WeChat
    print(f"\n[Step 5] Attach 微信进程 (PID={pid})")
    try:
        session = device.attach(pid)
        all_ok &= log_step("attach", True, f"attached to PID={pid}")
    except Exception as e:
        all_ok &= log_step("attach", False, str(e))
        save_report()
        return 1

    # Step 6: Load minimal key hunter script
    print("\n[Step 6] 加载 Key Hunter 脚本")
    js_code = """
    var captured = false;
    var keyInfo = null;

    function findModuleBase(moduleName) {
        var modules = Process.enumerateModules();
        for (var i = 0; i < modules.length; i++) {
            if (modules[i].name === moduleName) {
                return modules[i].base;
            }
        }
        return null;
    }

    function hookSqlite3Key() {
        var mod = Process.findModuleByName("libWCDB.so");
        if (!mod) {
            send({type: "info", payload: {msg: "libWCDB.so not loaded yet"}});
            return false;
        }
        var sqlite3Key = Module.findExportByName("libWCDB.so", "sqlite3_key");
        if (!sqlite3Key) {
            send({type: "error", payload: {msg: "sqlite3_key not found"}});
            return false;
        }
        send({type: "info", payload: {msg: "FOUND: sqlite3_key at " + sqlite3Key}});

        Interceptor.attach(sqlite3Key, {
            onEnter: function(args) {
                if (captured) return;
                var keyPtr = args[1];
                var keyLen = args[2].toInt32();
                if (keyLen > 0 && keyLen < 256) {
                    var key = Memory.readUtf8String(keyPtr, keyLen);
                    var keyHex = Memory.readByteArray(keyPtr, keyLen);
                    var hexStr = "";
                    var view = new Uint8Array(keyHex);
                    for (var i = 0; i < view.length; i++) {
                        hexStr += ("0" + view[i].toString(16)).slice(-2);
                    }
                    captured = true;
                    keyInfo = {key: key, hex: hexStr, len: keyLen};
                    send({type: "key", payload: {hex: hexStr, len: keyLen, via: "libWCDB.so!sqlite3_key"}});
                }
            }
        });
        return true;
    }

    // Try immediately, and retry if libWCDB.so not loaded
    if (!hookSqlite3Key()) {
        var retry = 0;
        var interval = setInterval(function() {
            retry++;
            if (hookSqlite3Key() || retry > 30) {
                clearInterval(interval);
                if (!captured && retry > 30) {
                    send({type: "error", payload: {msg: "Failed to hook sqlite3_key after 30 retries"}});
                }
            }
        }, 1000);
    }

    send({type: "info", payload: {msg: "Key hunter script loaded. Waiting for DB open..."}});
    """

    captured_keys = []
    def on_message(msg, data):
        if msg.get("type") != "send":
            return
        payload = msg.get("payload", {})
        ptype = payload.get("type", "")
        if ptype == "key":
            captured_keys.append(payload)
            print(f"       [KEY] hex={payload.get('hex')} len={payload.get('len')}")
        elif ptype == "info":
            print(f"       [INFO] {payload.get('msg', '')}")
        elif ptype == "error":
            print(f"       [ERROR] {payload.get('msg', '')}")

    try:
        script = session.create_script(js_code)
        script.on("message", on_message)
        script.load()
        all_ok &= log_step("script_load", True, "Key hunter loaded")
    except Exception as e:
        all_ok &= log_step("script_load", False, str(e))
        session.detach()
        save_report()
        return 1

    # Step 7: Wait for key capture
    print("\n[Step 7] 等待 key 捕获（最多 30 秒）...")
    print("       提示：打开微信，浏览聊天记录触发 DB 访问")
    for i in range(30):
        if captured_keys:
            break
        time.sleep(1)
        if i % 5 == 0:
            print(f"       ...已等待 {i} 秒")

    if captured_keys:
        key = captured_keys[0]
        all_ok &= log_step("key_capture", True, f"hex={key['hex'][:20]}... len={key['len']}")
        # Update config
        try:
            config = {}
            if os.path.exists("config.json"):
                with open("config.json") as f:
                    config = json.load(f)
            config["sqlcipher_key"] = bytes.fromhex(key["hex"]).decode("latin-1")
            with open("config.json", "w") as f:
                json.dump(config, f, indent=2, ensure_ascii=False)
            print(f"       [INFO] Key 已保存到 config.json")
        except Exception as e:
            print(f"       [WARN] 保存 key 失败: {e}")
    else:
        all_ok &= log_step("key_capture", False, "30秒内未捕获到 key")

    session.detach()

    # Summary
    REPORT["overall"] = "PASS" if all_ok else "FAIL"
    REPORT["timestamp"] = subprocess.check_output(["date", "-Iseconds"]).decode().strip()

    save_report()

    print("\n" + "=" * 60)
    print(f"结果: {REPORT['overall']}")
    print("=" * 60)
    return 0 if all_ok else 1

def save_report():
    with open("test_report_frida.json", "w") as f:
        json.dump(REPORT, f, indent=2, ensure_ascii=False)
    print("\n详细报告: test_report_frida.json")

if __name__ == "__main__":
    sys.exit(main())