127 lines
4.5 KiB
Python
127 lines
4.5 KiB
Python
# -*- 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()
|