gengx

wechat_rpa/test_wal_merge.py

@@ -0,0 +1,126 @@
+# -*- coding: utf-8 -*-
+"""
+test_wal_merge.py - 验证 WAL 帧合并方案（引擎 B 实时性增强）
+=============================================================
+思路：企微运行时新消息在 message.db-wal（未 checkpoint）。
+     WAL 帧数据与主库同算法加密（wxSQLite3 aes128cbc，帧头明文）。
+     本脚本：解密主库 → 解密 WAL 帧 → 按页号覆盖合并 → SQLite 打开验证。
+
+用法: python test_wal_merge.py <message.db> <wal_path> <key_hex>
+"""
+from __future__ import annotations
+
+import os
+import sqlite3
+import struct
+import sys
+import tempfile
+import time
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from wxwork_crypto import (
+    decrypt_db_to_file,
+    decrypt_page,
+    read_page_header,
+    SQLITE_FILE_HEADER,
+)
+
+
+def wal_read_frames(wal_path: str, key: bytes):
+    """读取 WAL 帧并解密，返回 {pgno: bytes}（仅 salt 匹配的帧，后帧覆盖前帧）。
+
+    返回 (pages, meta)：pages 为页号→解密页映射；meta 含 pgsz/salt/frame_count。
+    """
+    pages: dict[int, bytes] = {}
+    with open(wal_path, "rb") as f:
+        hdr = f.read(32)
+        if len(hdr) < 32:
+            return pages, {}
+        magic, ver, pgsz, ckpt, salt1, salt2, _, _ = struct.unpack(">IIIIIIII", hdr)
+        meta = {"pgsz": pgsz, "salt1": salt1, "salt2": salt2, "frames": 0, "applied": 0}
+        while True:
+            fh = f.read(24)
+            if len(fh) < 24:
+                break
+            pgno, commit, fs1, fs2, _, _ = struct.unpack(">IIIIII", fh)
+            data = f.read(pgsz)
+            if len(data) < pgsz:
+                break
+            meta["frames"] += 1
+            if fs1 != salt1 or fs2 != salt2:
+                continue  # 非当前 WAL 的帧（旧事务/已回滚），跳过
+            if pgno == 0 or pgno > (1 << 31):
+                continue
+            dec = decrypt_page(key, pgno, data, pgsz)
+            if dec is None:
+                continue
+            pages[pgno] = dec
+            meta["applied"] += 1
+    return pages, meta
+
+
+def merge_wal_into_db(db_path: str, wal_path: str, key: bytes, out_path: str) -> bool:
+    """解密主库 + 合并 WAL 帧 → out_path。返回是否成功。"""
+    # 1. 解密主库
+    if not decrypt_db_to_file(db_path, key, out_path):
+        return False
+    pages, meta = wal_read_frames(wal_path, key)
+    if not pages:
+        return True  # 无有效帧，主库副本即最终结果
+
+    pgsz = meta["pgsz"]
+    # 2. 按页号覆盖
+    with open(out_path, "r+b") as f:
+        size = os.path.getsize(out_path)
+        n_pages = (size + pgsz - 1) // pgsz
+        for pgno, data in pages.items():
+            if pgno > n_pages:
+                # 新页：扩展文件
+                f.seek(0, os.SEEK_END)
+                f.write(b"\x00" * ((pgno - n_pages) * pgsz))
+                n_pages = pgno
+            f.seek((pgno - 1) * pgsz)
+            f.write(data)
+        # 3. 更新 SQLite 头页数字段（offset 28，4 字节大端）
+        if n_pages > 0:
+            f.seek(28)
+            f.write(struct.pack(">I", n_pages))
+    return True
+
+
+def main():
+    if len(sys.argv) < 4:
+        print("用法: python test_wal_merge.py <message.db> <wal_path> <key_hex>")
+        sys.exit(1)
+    db_path, wal_path, key_hex = sys.argv[1], sys.argv[2], sys.argv[3]
+    key = bytes.fromhex(key_hex)
+
+    t0 = time.time()
+    tmp_fd, tmp_path = tempfile.mkstemp(suffix=".db", prefix="wal_merge_")
+    os.close(tmp_fd)
+    try:
+        ok = merge_wal_into_db(db_path, wal_path, key, tmp_path)
+        print(f"[+] 合并完成 ok={ok} 耗时 {time.time()-t0:.2f}s")
+        conn = sqlite3.connect(f"file:{tmp_path}?mode=ro", uri=True)
+        conn.text_factory = lambda b: b.decode("utf-8", errors="replace")
+        try:
+            cur = conn.cursor()
+            cur.execute("PRAGMA quick_check")
+            print("[+] quick_check:", cur.fetchone())
+            cur.execute("SELECT COUNT(*) FROM message_table")
+            print("[+] message_table 行数:", cur.fetchone()[0])
+            cur.execute("SELECT MAX(send_time) FROM message_table")
+            mx = cur.fetchone()[0]
+            print("[+] MAX(send_time):", mx,
+                  time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(mx)) if mx else "-")
+            # 与仅主库对比
+            cur.execute("SELECT MAX(send_time) FROM message_table WHERE send_time <= (SELECT MAX(send_time) FROM message_table)")
+        finally:
+            conn.close()
+    finally:
+        if os.path.exists(tmp_path):
+            os.remove(tmp_path)
+
+
+if __name__ == "__main__":
+    main()