gengx

wechat_rpa/test_crypto_roundtrip.py

@@ -0,0 +1,114 @@
+# -*- coding: utf-8 -*-
+"""
+test_crypto_roundtrip.py - wxwork_crypto 算法往返自测
+====================================================
+用 EncryptPage 逻辑（照抄 wxSQLite3 cipher_wxaes128.c）构造加密 db，
+再用 DecryptPage 逻辑解密验证。若往返一致，证明解密算法与官方实现对称，
+后续内存扫描失败的原因只可能是"密钥形态/位置"，而非算法错误。
+"""
+from __future__ import annotations
+
+import os
+import struct
+import sys
+import tempfile
+
+from Crypto.Cipher import AES
+
+from wxwork_crypto import (SQLITE_FILE_HEADER, MAGIC_SALT, decrypt_db_to_file,
+                           decrypt_page, generate_iv, page_key, read_page_header,
+                           verify_key)
+
+PAGE = 4096
+
+
+def _aes_cbc_encrypt(key, iv, data):
+    return AES.new(key, AES.MODE_CBC, iv).encrypt(data)
+
+
+def encrypt_page(master_key: bytes, page: int, page_data: bytes) -> bytes:
+    """EncryptPageAES128Cipher 非 legacy 逻辑。"""
+    data = bytearray(page_data)
+    key = page_key(master_key, page)
+    iv = generate_iv(page)
+    if page == 1:
+        db_header = bytes(data[16:24])
+        offset = 16
+        # 加密 0..15
+        data[0:16] = _aes_cbc_encrypt(key, iv, bytes(data[0:16]))
+        # 加密 16..end
+        data[offset:PAGE] = _aes_cbc_encrypt(key, iv, bytes(data[offset:PAGE]))
+        # 移动加密字节 16..23 → 8..15
+        data[8:16] = data[16:24]
+        # 恢复明文头
+        data[16:24] = db_header
+    else:
+        data[:] = _aes_cbc_encrypt(key, iv, bytes(data))
+    return bytes(data)
+
+
+def build_encrypted_db(master_key: bytes, n_pages: int = 8) -> bytes:
+    """构造一个假的加密 db：page1 标准 SQLite header + 其余页随机数据。"""
+    # 标准 SQLite header（100 字节）
+    header = bytearray(100)
+    header[0:16] = SQLITE_FILE_HEADER
+    header[16:18] = struct.pack(">H", PAGE)      # page size
+    header[18] = 2                                # write version (WAL)
+    header[19] = 2                                # read version
+    header[20] = 0                                # reserved
+    header[21] = 0x40                             # max payload
+    header[22] = 0x20                             # min payload
+    header[23] = 0x20                             # leaf payload
+    header[28:32] = struct.pack(">I", n_pages)    # 页数
+    header[96:100] = b"\x00\x00\x00\x00"          # change counter 占位
+    page1 = bytes(header) + b"\x00" * (PAGE - 100)
+
+    out = bytearray()
+    for pg in range(1, n_pages + 1):
+        if pg == 1:
+            raw = page1
+        else:
+            raw = bytes((pg * 7 + i * 13) & 0xFF for i in range(PAGE))  # 伪随机页
+        out += encrypt_page(master_key, pg, raw)
+    return bytes(out)
+
+
+def main():
+    ok = True
+    for key_len in (16, 32):
+        master_key = bytes(range(1, key_len + 1))
+        blob = build_encrypted_db(master_key, n_pages=8)
+        fd, path = tempfile.mkstemp(suffix=".db")
+        os.close(fd)
+        with open(path, "wb") as f:
+            f.write(blob)
+        try:
+            info = read_page_header(path)
+            assert info is not None, "read_page_header 应识别加密头"
+            assert info[0] == PAGE, f"页大小 {info[0]} != {PAGE}"
+            assert verify_key(master_key, path), "verify_key 应通过"
+            # 错误密钥应失败
+            bad = bytes(0xEE for _ in range(key_len))
+            assert not verify_key(bad, path), "错误密钥不应通过"
+
+            out_path = path + ".dec"
+            ok2 = decrypt_db_to_file(path, master_key, out_path)
+            assert ok2, "整库解密应成功"
+            with open(out_path, "rb") as f:
+                dec = f.read()
+            assert dec[0:16] == SQLITE_FILE_HEADER, "解密后 magic 应正确"
+            assert dec[16:18] == struct.pack(">H", PAGE), "页大小头应还原"
+            assert dec[18] == 2 and dec[19] == 2, "版本字节应还原"
+            assert dec[21] == 0x40 and dec[22] == 0x20 and dec[23] == 0x20
+            print(f"[PASS] key_len={key_len}: 往返解密正确, 页1头还原完整")
+        finally:
+            for p in (path, path + ".dec"):
+                if os.path.exists(p):
+                    os.remove(p)
+    if ok:
+        print("test_crypto_roundtrip: ALL PASS")
+    return 0 if ok else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())