115 lines
4.3 KiB
Python
115 lines
4.3 KiB
Python
# -*- 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())
|