WechatOnCloud/bridge/woc_bridge/db/decryptor.py

"""纯 Python 微信 4.x Linux SQLCipher/WCDB 解密器。

参考 wechat-cli-main 项目实现，针对微信 4.x Linux 加密数据库：
- 页面大小 4096 字节
- 第 1 页前 16 字节为 salt
- 保留区 80 字节：IV(16) + HMAC-SHA512(64)
- AES-256-CBC 解密
- HMAC-SHA512 页校验
- mac_salt = salt XOR 0x3A，PBKDF2-HMAC-SHA512 2 轮派生 mac_key
"""

from __future__ import annotations

import hashlib
import hmac
import os
import shutil
import struct

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes


# SQLite 文件头
SQLITE_HEADER = b"SQLite format 3\x00"
PAGE_SIZE = 4096
KEY_SIZE = 32
SALT_SIZE = 16
IV_SIZE = 16
HMAC_SIZE = 64  # SHA512 HMAC
RESERVE_SIZE = 80  # IV(16) + HMAC(64)
ROUND_COUNT = 256000  # PBKDF2 迭代轮数（微信 4.x Linux 使用 SHA512）
MAC_SALT_XOR = 0x3A  # mac_salt = salt ^ 0x3A

# WAL 文件常量
WAL_HEADER_SIZE = 32
WAL_FRAME_HEADER_SIZE = 24


def _derive_enc_key(key_material: bytes, salt: bytes) -> bytes:
    """从 32 字节 key material 派生 AES 加密密钥（PBKDF2-SHA512 256000 轮）。"""
    return hashlib.pbkdf2_hmac("sha512", key_material, salt, ROUND_COUNT, dklen=KEY_SIZE)


def _derive_mac_key(enc_key: bytes, salt: bytes) -> bytes:
    """派生微信 4.x 页 HMAC 密钥（mac_salt = salt XOR 0x3A，PBKDF2-SHA512 2 轮）。"""
    mac_salt = bytes(b ^ MAC_SALT_XOR for b in salt)
    return hashlib.pbkdf2_hmac("sha512", enc_key, mac_salt, 2, dklen=KEY_SIZE)


def _compute_page_hmac(mac_key: bytes, page: bytes, page_num: int) -> bytes:
    """计算页面 HMAC-SHA512。

    微信 4.x Linux 页 HMAC 计算范围（含 IV）：
    - 第 1 页：salt 之后到 HMAC 之前的数据（即 offset=16 到 data_end=4096-80+16=4032）
    - 其他页：从头开始到 HMAC 之前的数据（offset=0 到 data_end=4032）
    最后追加 4 字节小端页号。
    """
    offset = SALT_SIZE if page_num == 1 else 0
    data_end = PAGE_SIZE - RESERVE_SIZE + IV_SIZE
    mac = hmac.new(mac_key, digestmod=hashlib.sha512)
    mac.update(page[offset:data_end])
    mac.update(page_num.to_bytes(4, "little"))
    return mac.digest()


def _resolve_page1_key_material(
    key_material: bytes, page1: bytes
) -> tuple[bytes, bytes, str] | None:
    """检测输入密钥是否匹配第 1 页 HMAC。

    先尝试把输入当作已派生的 enc_key 直接验证；失败则当作 raw key_material
    派生 enc_key 后再验证。返回 (enc_key, mac_key, mode) 或 None。
    """
    if len(page1) < PAGE_SIZE:
        return None

    salt = page1[:SALT_SIZE]
    stored_page1_hmac = page1[PAGE_SIZE - HMAC_SIZE: PAGE_SIZE]

    # 尝试 1：输入就是 enc_key
    mac_key = _derive_mac_key(key_material, salt)
    if hmac.compare_digest(stored_page1_hmac, _compute_page_hmac(mac_key, page1, 1)):
        return key_material, mac_key, "enc_key"

    # 尝试 2：输入是 raw key_material，需派生 enc_key
    enc_key = _derive_enc_key(key_material, salt)
    mac_key = _derive_mac_key(enc_key, salt)
    if hmac.compare_digest(stored_page1_hmac, _compute_page_hmac(mac_key, page1, 1)):
        return enc_key, mac_key, "key_material"

    return None


def _decrypt_page(enc_key: bytes, page: bytes, page_num: int) -> bytes:
    """使用 AES-256-CBC 解密单页。

    明文 SQLite 页不携带加密保留区，通过零填充保留页大小不变。
    """
    iv = page[PAGE_SIZE - RESERVE_SIZE: PAGE_SIZE - RESERVE_SIZE + IV_SIZE]
    offset = SALT_SIZE if page_num == 1 else 0
    encrypted_page = page[offset: PAGE_SIZE - RESERVE_SIZE]

    cipher = Cipher(
        algorithms.AES(enc_key),
        modes.CBC(iv),
        backend=default_backend(),
    )
    decryptor = cipher.decryptor()
    decrypted_page = decryptor.update(encrypted_page) + decryptor.finalize()

    if page_num == 1:
        return SQLITE_HEADER + decrypted_page + (b"\x00" * RESERVE_SIZE)
    return decrypted_page + (b"\x00" * RESERVE_SIZE)


def decrypt_wal(enc_key: bytes, wal_path: str, output_db_path: str) -> int:
    """把加密的 WAL 文件合并到已解密的明文 DB 中。

    微信 4.x Linux 使用 SQLCipher WAL 格式，每帧结构为：
    - frame header (24 bytes): pgno(4) + commit_size(4) + salt1(4) + salt2(4) + checksum(8)
    - frame page (4096 bytes): 加密页数据

    只合并与 WAL header 中 salt 匹配的帧，避免应用不完整数据。

    Args:
        enc_key: 32 字节 AES 加密密钥
        wal_path: 加密 WAL 文件路径
        output_db_path: 已解密的明文 DB 路径（会被原地 patch）

    Returns:
        int: 成功合并的帧数
    """
    if not os.path.exists(wal_path):
        return 0

    wal_size = os.path.getsize(wal_path)
    if wal_size <= WAL_HEADER_SIZE:
        return 0

    patched = 0
    with open(wal_path, "rb") as wf, open(output_db_path, "r+b") as df:
        wal_header = wf.read(WAL_HEADER_SIZE)
        if len(wal_header) < WAL_HEADER_SIZE:
            return 0

        wal_salt1 = struct.unpack(">I", wal_header[16:20])[0]
        wal_salt2 = struct.unpack(">I", wal_header[20:24])[0]
        frame_size = WAL_FRAME_HEADER_SIZE + PAGE_SIZE

        while wf.tell() + frame_size <= wal_size:
            frame_header = wf.read(WAL_FRAME_HEADER_SIZE)
            if len(frame_header) < WAL_FRAME_HEADER_SIZE:
                break

            pgno = struct.unpack(">I", frame_header[0:4])[0]
            frame_salt1 = struct.unpack(">I", frame_header[8:12])[0]
            frame_salt2 = struct.unpack(">I", frame_header[12:16])[0]
            encrypted_page = wf.read(PAGE_SIZE)
            if len(encrypted_page) < PAGE_SIZE:
                break

            # 过滤异常帧
            if pgno == 0 or pgno > 10000000:
                continue
            # salt 不匹配说明 WAL 已轮换或数据不完整，跳过
            if frame_salt1 != wal_salt1 or frame_salt2 != wal_salt2:
                continue

            decrypted_page = _decrypt_page(enc_key, encrypted_page, pgno)
            df.seek((pgno - 1) * PAGE_SIZE)
            df.write(decrypted_page)
            patched += 1

    return patched


class Decryptor:
    """微信 4.x Linux SQLCipher/WCDB 数据库解密器。"""

    def __init__(self, key_hex: str):
        """初始化解密器，key_hex 为 64 位十六进制密钥。

        Args:
            key_hex: 64 位十六进制字符串（32 字节 enc_key 或 raw key_material）

        Raises:
            ValueError: 密钥长度不为 64 或非有效十六进制
        """
        if len(key_hex) != 64:
            raise ValueError("密钥必须是 64 位十六进制字符串")

        try:
            self.key_bytes = bytes.fromhex(key_hex)
        except ValueError:
            raise ValueError("密钥必须是有效的十六进制字符串")

        self._key_hex = key_hex
        self._enc_key: bytes | None = None
        self._mac_key: bytes | None = None
        self._key_mode: str = ""

    def _resolve(self, page1: bytes) -> bool:
        """用第 1 页解析出 enc_key / mac_key / mode，成功返回 True。"""
        resolved = _resolve_page1_key_material(self.key_bytes, page1)
        if resolved is None:
            return False
        self._enc_key, self._mac_key, self._key_mode = resolved
        return True

    def decrypt_db(self, enc_db_path: str, output_path: str) -> dict:
        """解密微信 4.x 加密数据库到明文 SQLite 文件。

        采用流式分页读写：
        - 仅把当前页保留在内存，不一次性读入整库。
        - 明文 SQLite 直接流式复制。

        Args:
            enc_db_path: 加密 DB 文件路径
            output_path: 解密后明文 SQLite 文件路径

        Returns:
            dict: 解密结果统计信息
        """
        result: dict = {
            "success": False,
            "key_mode": "",
            "total_pages": 0,
            "successful_pages": 0,
            "failed_pages": 0,
            "copied_as_sqlite": False,
            "error": "",
        }

        try:
            file_size = os.path.getsize(enc_db_path)
        except Exception as e:
            result["error"] = f"stat_error: {e}"
            return result

        if file_size < PAGE_SIZE:
            result["error"] = "file_too_small"
            return result

        try:
            with open(enc_db_path, "rb") as f:
                page1 = f.read(PAGE_SIZE)
        except Exception as e:
            result["error"] = f"read_error: {e}"
            return result

        if page1.startswith(SQLITE_HEADER):
            try:
                with open(enc_db_path, "rb") as src, open(output_path, "wb") as dst:
                    shutil.copyfileobj(src, dst)
            except Exception as e:
                result["error"] = f"write_error: {e}"
                return result
            result["success"] = True
            result["copied_as_sqlite"] = True
            return result

        if not self._resolve(page1):
            result["error"] = "key_mismatch"
            return result

        result["key_mode"] = self._key_mode

        total_pages = (file_size + PAGE_SIZE - 1) // PAGE_SIZE
        result["total_pages"] = total_pages
        successful_pages = 0
        failed_pages = 0

        try:
            with open(enc_db_path, "rb") as src, open(output_path, "wb") as dst:
                for page_num in range(1, total_pages + 1):
                    page = src.read(PAGE_SIZE)
                    if not page:
                        break
                    if len(page) < PAGE_SIZE:
                        page = page + (b"\x00" * (PAGE_SIZE - len(page)))

                    # 直接解密，不做逐页 HMAC 验证
                    # 参考 wechat-cli-main：微信运行时 DB 页面可能因 WAL 并发写入
                    # 等原因 HMAC 不匹配，但 key 本身正确（page1 验证已通过）。
                    # 逐页 HMAC 验证导致大量页面误判为损坏，实际 AES 解密仍可得到有效数据
                    try:
                        dst.write(_decrypt_page(self._enc_key, page, page_num))
                        successful_pages += 1
                    except Exception:
                        failed_pages += 1
                        dst.write(b"\x00" * PAGE_SIZE)
        except Exception as e:
            result["error"] = f"write_error: {e}"
            return result

        result["successful_pages"] = successful_pages
        result["failed_pages"] = failed_pages
        result["success"] = True
        return result

    def verify_key(self, enc_db_path: str) -> bool:
        """便捷方法：验证 key 是否匹配指定 DB 文件。"""
        try:
            with open(enc_db_path, "rb") as f:
                data = f.read(PAGE_SIZE)
        except Exception:
            return False

        if len(data) < PAGE_SIZE:
            return False

        if data.startswith(SQLITE_HEADER):
            return True

        return self._resolve(data)

    def decrypt_wal(self, wal_path: str, output_db_path: str) -> int:
        """使用已解析的 enc_key 合并 WAL 到明文 DB。

        必须先调用 decrypt_db 或 verify_key 解析出 enc_key。
        """
        if self._enc_key is None:
            raise RuntimeError("enc_key 尚未解析，请先调用 decrypt_db 或 verify_key")
        return decrypt_wal(self._enc_key, wal_path, output_db_path)