gengx

wechat_rpa/test_scan_probe_mis.py

@@ -0,0 +1,88 @@
+# -*- 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())