# -*- coding: utf-8 -*-
"""
test_scan_probe_mis.py - 验证全滑窗对"非 16 字节对齐"密钥的命中能力
==================================================================
企微堆上密钥未必 16 字节对齐。本探针把密钥放在 5 字节偏移处，
先 aligned 扫（预期可能不命中），再全滑窗扫（预期命中）。
同时测量全滑窗在真实企微场景下的耗时量级。
"""
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_mis(key_hex: str, db_path: str, duration: int = 120) -> subprocess.Popen:
    code = f"""
import time, sys
key = bytes.fromhex({key_hex!r})
db = {db_path!r}
# 非对齐: 前导 5 字节 -> key 起始偏移 5+64=69, 69 % 16 = 5, 非对齐
holder = [key, b'\\x00' * 5 + key + b'\\x00' * 32]
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_mis(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 失败"

    t0 = time.time()
    keys_a = wk.scan_binary_single(pid, prep, aligned=True, progress_cb=lambda m: None)
    dt_a = time.time() - t0
    print(f"[3/4] aligned 扫描 {dt_a:.1f}s, 命中 {len(keys_a)} 个 "
          f"({'✅' if key16.hex() in keys_a else '❌(预期可能不命中,非对齐)'})")

    if key16.hex() not in keys_a:
        t0 = time.time()
        keys_f = wk.scan_binary_single(pid, prep, aligned=False, progress_cb=lambda m: None)
        dt_f = time.time() - t0
        print(f"[4/4] full-slide 扫描 {dt_f:.1f}s, 命中 {len(keys_f)} 个 "
              f"({'✅ 全滑窗可兜底' if key16.hex() in keys_f else '❌ 全滑窗也有盲区'})")
        hit = key16.hex() in keys_f
    else:
        print("[4/4] aligned 已命中, 跳过全滑窗")
        hit = True

    proc.terminate()
    proc.wait()
    os.remove(db_path)
    return 0 if hit else 1


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