# -*- 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())