# -*- coding: utf-8 -*-
"""
wxwork_crypto.py - 企业微信(WXWork)本地数据库解密算法核心
=============================================================
对应 wxSQLite3 "aes128cbc" cipher（SQLite3MultipleCiphers 项目 cipher_wxaes128.c / codec_algos.c）
企微 5.x 使用的加密方案：
  - 算法    : AES-128-CBC（无 HMAC、无每页保留字节）
  - 主密钥  : 16 字节（从 WXWork.exe 进程内存提取，或经 PDF 规范 MD5+RC4 派生）
  - 页面密钥: MD5(master_key || page_le32 || "sAlT") 取 16 字节
  - 每页 IV : 线性同余(页号) → 16 字节 initkey → MD5
  - 页 1 布局: 0..15 加密(不解密) / 16..23 明文头 / 加密字节暂存于 8..15 / 24..end 加密

参考实现:
  https://github.com/utelle/SQLite3MultipleCiphers
    src/cipher_wxaes128.c   (GenerateKey / EncryptPage / DecryptPage)
    src/codec_algos.c       (sqlite3mcAES128 / sqlite3mcGenerateInitialVector / RC4 / MD5)

依赖: pycryptodome (AES), hashlib (MD5)
"""
from __future__ import annotations

import hashlib
import os
import sqlite3
import struct
import tempfile

from Crypto.Cipher import AES

SQLITE_FILE_HEADER = b"SQLite format 3\x00"
MAGIC_SALT = b"sAlT"  # 0x73 0x41 0x6c 0x54
PAGE1_HEADER_CHECK = (0x40, 0x20, 0x20)  # dbHeader[5..7] 魔数校验


def generate_iv(page: int) -> bytes:
    """每页 IV：MODMULT 线性同余 4 轮生成 16 字节 initkey，再 MD5。"""
    z = page + 1
    initkey = bytearray(16)
    for j in range(4):
        # MODMULT(52774, 40692, 3791, 2147483399, s)
        q = z // 52774
        z = 40692 * (z - 52774 * q) - 3791 * q
        if z < 0:
            z += 2147483399
        initkey[4 * j + 0] = 0xFF & z
        initkey[4 * j + 1] = 0xFF & (z >> 8)
        initkey[4 * j + 2] = 0xFF & (z >> 16)
        initkey[4 * j + 3] = 0xFF & (z >> 24)
    return hashlib.md5(bytes(initkey)).digest()


def page_key(master_key: bytes, page: int) -> bytes:
    """页面密钥：MD5/SHA256(master_key || page_le32 || 'sAlT')。

    - 16 字节主密钥(waAES128) → MD5 → 16 字节，AES-128-CBC
    - 32 字节主密钥(wxAES256) → SHA256 → 32 字节，AES-256-CBC
    """
    nkey = master_key + struct.pack("<I", page & 0xFFFFFFFF) + MAGIC_SALT
    if len(master_key) == 16:
        return hashlib.md5(nkey).digest()
    return hashlib.sha256(nkey).digest()


def _aes_cbc_decrypt(key16: bytes, iv16: bytes, data: bytes) -> bytes:
    return AES.new(key16, AES.MODE_CBC, iv16).decrypt(data)


def read_page_header(db_path: str) -> tuple | None:
    """读取 db 文件第 1 页明文头 16..23，返回 (page_size, header8) 或 None(文件太小/无 db)。"""
    try:
        with open(db_path, "rb") as f:
            head = f.read(24)
        if len(head) < 24:
            return None
        header8 = head[16:24]
        db_page_size = (header8[0] << 8) | (header8[1] << 16)
        if db_page_size < 512 or ((db_page_size - 1) & db_page_size) != 0:
            return None
        if header8[5] != PAGE1_HEADER_CHECK[0] or header8[6] != PAGE1_HEADER_CHECK[1] or header8[7] != PAGE1_HEADER_CHECK[2]:
            return None
        return db_page_size, header8
    except Exception:
        return None


def decrypt_page1_block(master_key: bytes, db_path: str) -> bytes | None:
    """仅解密第 1 页的 16..32 密文块，用于快速校验密钥。

    加密文件里 page1 布局: 0..15 加密 / 8..15 暂存加密的版本头 /
    16..23 明文头 / 24..end 加密。解密数据流需拼接 8..16 + 24..32，
    即 data[16:32] = (8..16 的密文A) + (24..32 的密文B)。
    返回解密后的 16 字节明文（前 8 字节应为明文头 dbHeader）或 None。
    支持 16 字节(AES-128) 与 32 字节(AES-256) 主密钥。
    """
    if len(master_key) not in (16, 32):
        return None
    info = read_page_header(db_path)
    if info is None:
        return None
    page_size, header8 = info
    with open(db_path, "rb") as f:
        f.seek(8)
        part_a = f.read(8)   # 8..16: 加密的版本头
        f.seek(24)
        part_b = f.read(8)   # 24..32: 加密的后续数据
    if len(part_a) < 8 or len(part_b) < 8:
        return None
    block = part_a + part_b  # 对应解密数据流 data[16:32]
    key = page_key(master_key, 1)
    iv = generate_iv(1)
    return _aes_cbc_decrypt(key, iv, block)


def verify_key(master_key: bytes, db_path: str) -> bool:
    """校验主密钥是否能解密该 db 文件（页1明文头匹配）。支持 16/32 字节。"""
    if len(master_key) not in (16, 32):
        return False
    info = read_page_header(db_path)
    if info is None:
        return False
    _, header8 = info
    dec = decrypt_page1_block(master_key, db_path)
    if dec is None or len(dec) < 8:
        return False
    return dec[:8] == header8


def decrypt_page(master_key: bytes, page: int, page_data: bytes, page_size: int) -> bytes | None:
    """解密单个页面（非 legacy 模式）。

    - page 1：16..23 为明文头，加密字节从 8..15 恢复；0..15 不解密，覆写为 SQLite magic。
    - 其他页：整页解密。
    密钥错误时返回 None（仅页 1 能自校验；其他页需由 SQLite 解析兜底）。
    """
    data = bytearray(page_data)
    key = page_key(master_key, page)
    iv = generate_iv(page)
    if page == 1:
        header8 = bytes(data[16:24])
        # 校验明文头魔数，避免对非本格式文件瞎解
        if header8[5] != PAGE1_HEADER_CHECK[0] or header8[6] != PAGE1_HEADER_CHECK[1] or header8[7] != PAGE1_HEADER_CHECK[2]:
            return None
        data[16:24] = data[8:16]  # 恢复加密字节
        dec = _aes_cbc_decrypt(key, iv, bytes(data[16:page_size]))
        data[16:page_size] = dec
        if bytes(data[16:24]) != header8:
            return None  # 密钥错误
        data[0:16] = SQLITE_FILE_HEADER
        return bytes(data)
    else:
        dec = _aes_cbc_decrypt(key, iv, bytes(data[:page_size]))
        return bytes(dec)


def decrypt_db_to_file(db_path: str, master_key: bytes, out_path: str, page_size: int = 4096,
                       progress_cb=None) -> bool:
    """整库解密为标准 SQLite 文件。

    成功返回 True；密钥错误/文件异常返回 False。
    """
    info = read_page_header(db_path)
    if info is None:
        return False
    db_page_size, _ = info
    size = os.path.getsize(db_path)
    if size == 0:
        return False
    n_pages = (size + db_page_size - 1) // db_page_size
    with open(db_path, "rb") as fin, open(out_path, "wb") as fout:
        for pg in range(1, n_pages + 1):
            data = fin.read(db_page_size)
            if len(data) < db_page_size:
                data += b"\x00" * (db_page_size - len(data))
            dec = decrypt_page(master_key, pg, data, db_page_size)
            if dec is None:
                return False
            fout.write(dec)
            if progress_cb and pg % 256 == 0:
                progress_cb(pg, n_pages)
    return True


def decrypt_db_open(db_path: str, master_key: bytes) -> sqlite3.Connection | None:
    """解密 db 到临时文件并用 sqlite3 打开（只读）。调用方负责 close 并删除临时文件。

    返回 (conn, tmp_path) 或 (None, None)。
    """
    tmp_fd, tmp_path = tempfile.mkstemp(suffix=".db", prefix="wxwork_dec_")
    os.close(tmp_fd)
    try:
        ok = decrypt_db_to_file(db_path, master_key, tmp_path)
        if not ok:
            os.remove(tmp_path)
            return None, None
        conn = sqlite3.connect(f"file:{tmp_path}?mode=ro", uri=True)
        return conn, tmp_path
    except Exception:
        if os.path.exists(tmp_path):
            try:
                os.remove(tmp_path)
            except Exception:
                pass
        return None, None


# ---------------------------------------------------------------------------
# PDF 规范密钥派生（备用：内存中找到的是 passphrase 而非 raw key 时使用）
# ---------------------------------------------------------------------------
_PDF_PAD = b"\x28" * 32  # PDF 规范口令填充


def _pad_password(passphrase: bytes) -> bytes:
    if len(passphrase) >= 32:
        return passphrase[:32]
    return passphrase + _PDF_PAD[: 32 - len(passphrase)]


def _rc4(key: bytes, data: bytes) -> bytes:
    S = list(range(256))
    j = 0
    klen = len(key)
    for i in range(256):
        j = (j + S[i] + key[i % klen]) % 256
        S[i], S[j] = S[j], S[i]
    out = bytearray(len(data))
    a = b = 0
    for i in range(len(data)):
        a = (a + 1) % 256
        t = S[a]
        b = (b + t) % 256
        S[a], S[b] = S[b], S[a]
        out[i] = data[i] ^ S[(S[a] + S[b]) % 256]
    return bytes(out)


def derive_key_pdf(passphrase: bytes) -> bytes:
    """GenerateKeyAES128Cipher 的完整实现（PDF 规范 MD5+RC4），返回 16 字节 AES key。"""
    user_pad = _pad_password(passphrase)
    owner_pad = _pad_password(b"")

    # --- owner key ---
    digest = hashlib.md5(owner_pad).digest()
    for _ in range(50):
        digest = hashlib.md5(digest).digest()
    owner_key = bytearray(user_pad)
    for i in range(20):
        mkey = bytes((digest[j] ^ i) & 0xFF for j in range(16))
        owner_key = bytearray(_rc4(mkey, bytes(owner_key)))

    # --- encryption key ---
    digest = hashlib.md5(user_pad + bytes(owner_key)).digest()
    for _ in range(50):
        digest = hashlib.md5(digest).digest()
    return digest[:16]


if __name__ == "__main__":
    # 自检：算法一致性（生成IV/页面密钥长度/派生key长度）
    iv1 = generate_iv(1)
    iv2 = generate_iv(2)
    assert len(iv1) == 16 and len(iv2) == 16 and iv1 != iv2
    k = page_key(b"\x01" * 16, 1)
    assert len(k) == 16
    dk = derive_key_pdf(b"test")
    assert len(dk) == 16
    print("wxwork_crypto self-check OK")