# -*- coding: utf-8 -*-
"""
test_scan_probe.py - 验证密钥扫描器对真实内存中密钥的命中能力
============================================================
启动一个子进程，其内存中放置已知 16/32 字节密钥（含真实加密 db 验证），
再用扫描器扫描该子进程，确认能命中。通过则证明扫描链路无盲区，
企微密钥确实不在可读内存（需重启企微黄金窗口提取）。
"""
from __future__ import annotations

import os
import subprocess
import sys
import tempfile
import time

from test_crypto_roundtrip import build_encrypted_db  # 复用假加密库构造


def _spawn_probe(key_hex: str, db_path: str, duration: int = 90) -> subprocess.Popen:
    # 用 !r 生成合法 Python 字面量，避免 Windows 路径中的 \U \x 等被误解析为转义
    code = f"""
import time, sys
key = bytes.fromhex({key_hex!r})
db = {db_path!r}
# 模拟企微: 密钥在堆上（列表持有引用）+ 打开 db 句柄
holder = [key, key, b'\\x00' * 64 + key + b'\\x00' * 64]
try:
    f = open(db, 'rb')
except Exception:
    f = None
t_end = time.time() + {duration}
while time.time() < t_end:
    time.sleep(0.2)
"""
    return subprocess.Popen([sys.executable, "-c", code])


def main():
    import multiprocessing as mp
    mp.freeze_support()
    print("[1/4] 构造假加密 db ...")
    key16 = bytes(range(1, 17))
    blob = build_encrypted_db(key16, n_pages=4)
    fd, db_path = tempfile.mkstemp(suffix=".db")
    os.close(fd)
    with open(db_path, "wb") as f:
        f.write(blob)

    print("[2/4] 启动携带密钥的探针子进程 ...")
    proc = _spawn_probe(key16.hex(), db_path)
    time.sleep(3)  # 等子进程就绪
    if proc.poll() is not None:
        print(f"      ❌ 探针子进程提前退出 rc={proc.returncode}，无法扫描")
        os.remove(db_path)
        return 1
    pid = proc.pid

    import wxwork_key as wk
    from wxwork_crypto import generate_iv
    wk._IV_PAGE1 = generate_iv(1)
    prep = wk._prep_verify(db_path)
    assert prep is not None, "prep 失败"

    print(f"[3/4] 扫描探针进程 PID={pid} ...")
    t0 = time.time()
    # 先对齐扫描（快 16 倍）；未命中再全滑窗兜底
    keys = wk.scan_binary_single(pid, prep, aligned=True, progress_cb=lambda m: None)
    mode = "aligned"
    if not keys:
        keys = wk.scan_binary_single(pid, prep, aligned=False, progress_cb=lambda m: None)
        mode = "full-slide"
    print(f"      {mode} 扫描耗时 {time.time()-t0:.1f}s, 命中 {len(keys)} 个")
    hit = key16.hex() in keys
    print(f"[4/4] 结果: {'命中 ✅ 扫描链路正常' if hit else '未命中 ❌ 扫描器有盲区'}")
    proc.terminate()
    proc.wait()
    os.remove(db_path)
    return 0 if hit else 1


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