"""SQLCipher 4 raw-key reader. Native AES; authenticated pages and WAL commits.

Only 4096-byte pages / 80 reserved bytes / SHA512 / raw AES keys are supported.
No source database is opened for writing. Unknown formats fail closed.
"""
import ctypes as C
import hashlib
import hmac
import os
import platform
import sqlite3
import struct
from functools import lru_cache
from pathlib import Path

PAGE = 4096
RESERVE = 80


@lru_cache(maxsize=1)
def common_crypto():
    lib = C.CDLL('/usr/lib/system/libcommonCrypto.dylib')
    fn = lib.CCCrypt
    fn.argtypes = [C.c_uint, C.c_uint, C.c_uint, C.c_void_p, C.c_size_t,
                   C.c_void_p, C.c_void_p, C.c_size_t, C.c_void_p, C.c_size_t, C.POINTER(C.c_size_t)]
    return fn


def aes_cbc(data, key, iv, encrypt=False):
    if len(key) != 32 or len(iv) != 16 or len(data) % 16:
        raise ValueError('无效的 AES 数据长度')
    output = C.create_string_buffer(len(data) + 16)
    if platform.system() == 'Darwin':
        fn = common_crypto()
        moved = C.c_size_t()
        code = fn(0 if encrypt else 1, 0, 0, key, len(key), iv, data, len(data), output, len(output), C.byref(moved))
        if code or moved.value != len(data):
            raise ValueError('本机 AES 解密失败')
        return output.raw[:moved.value]
    if platform.system() == 'Windows':
        # Windows native crypto also permits fixture tests without pip packages.
        lib = C.WinDLL('bcrypt')
        handle, key_handle = C.c_void_p(), C.c_void_p()
        def check(result):
            if result:
                raise ValueError('本机 AES 操作失败')
        check(lib.BCryptOpenAlgorithmProvider(C.byref(handle), C.c_wchar_p('AES'), None, 0))
        try:
            mode = C.create_unicode_buffer('ChainingModeCBC')
            check(lib.BCryptSetProperty(handle, C.c_wchar_p('ChainingMode'), mode, C.sizeof(mode), 0))
            length, written = C.c_ulong(), C.c_ulong()
            check(lib.BCryptGetProperty(handle, C.c_wchar_p('ObjectLength'), C.byref(length), 4, C.byref(written), 0))
            obj = C.create_string_buffer(length.value)
            check(lib.BCryptGenerateSymmetricKey(handle, C.byref(key_handle), obj, length.value, key, len(key), 0))
            ivbuf = C.create_string_buffer(iv, 16)
            fn = lib.BCryptEncrypt if encrypt else lib.BCryptDecrypt
            check(fn(key_handle, data, len(data), None, ivbuf, 16, output, len(output), C.byref(written), 0))
            return output.raw[:written.value]
        finally:
            if key_handle.value:
                lib.BCryptDestroyKey(key_handle)
            lib.BCryptCloseAlgorithmProvider(handle, 0)
    raise ValueError('原生解密目前支持 macOS / Windows')


def hmac_key(key, salt):
    return hashlib.pbkdf2_hmac('sha512', key, bytes(x ^ 0x3a for x in salt), 2, 32)


def decrypt_page(page, number, key, mac_key):
    if len(page) != PAGE:
        raise ValueError('数据库页大小不匹配（仅支持 SQLCipher 4 / 4096）')
    if not any(page) and number != 1:
        return page
    start = 16 if number == 1 else 0
    payload, iv = page[start:-RESERVE], page[-RESERVE:-64]
    expected = hmac.new(mac_key, payload + iv + struct.pack('<I', number), hashlib.sha512).digest()
    if not hmac.compare_digest(expected, page[-64:]):
        raise ValueError('页认证失败：密钥、格式不匹配或副本损坏；没有跳过错误页')
    plain = aes_cbc(payload, key, iv)
    return (b'SQLite format 3\x00' if number == 1 else b'') + plain + bytes(RESERVE)


def matches(first_page, key):
    try:
        plain = decrypt_page(first_page, 1, key, hmac_key(key, first_page[:16]))
        return plain[16:18] == b'\x10\x00' and plain[20] == RESERVE
    except ValueError:
        return False


def wal_checksum(data, initial=(0, 0), endian='<'):
    if len(data) % 8:
        raise ValueError('WAL 校验数据长度错误')
    a, b = initial
    for x, y in struct.iter_unpack(endian + 'II', data):
        a = (a + x + b) & 0xffffffff
        b = (b + y + a) & 0xffffffff
    return a, b


def wal_frames(path):
    """Validate the complete WAL chain; return only frames through last commit."""
    if not path.exists() or path.stat().st_size == 0:
        return [], None
    frames, commit, db_pages = [], 0, None
    with path.open('rb') as file:
        header = file.read(32)
        if len(header) != 32:
            raise ValueError('WAL 头不完整，请退出微信后重新备份')
        magic, version, size = struct.unpack('>III', header[:12])
        if magic not in (0x377f0682, 0x377f0683) or version != 3007000 or size != PAGE:
            raise ValueError('不支持的 WAL 格式')
        endian = '<' if magic == 0x377f0682 else '>'
        checksum = wal_checksum(header[:24], endian=endian)
        if checksum != struct.unpack('>II', header[24:32]):
            raise ValueError('WAL 头校验失败')
        while True:
            offset = file.tell()
            frame_header = file.read(24)
            if not frame_header:
                break
            page = file.read(PAGE)
            if len(frame_header) != 24 or len(page) != PAGE:
                raise ValueError('WAL 尾帧不完整，请重新取得稳定副本')
            if frame_header[8:16] != header[16:24]:
                raise ValueError('WAL salt 不匹配，副本可能跨越了数据库轮换')
            checksum = wal_checksum(frame_header[:8] + page, checksum, endian)
            if checksum != struct.unpack('>II', frame_header[16:24]):
                raise ValueError('WAL 帧校验失败，请重新备份')
            number, count = struct.unpack('>II', frame_header[:8])
            if not number or number > 10_000_000 or count > 10_000_000:
                raise ValueError('WAL 页编号无效')
            frames.append((offset + 24, number))
            if count:
                commit, db_pages = len(frames), count
    return frames[:commit], db_pages


def decrypt_database(source, destination, key):
    source, destination = Path(source), Path(destination)
    with source.open('rb') as file:
        first = file.read(PAGE)
    mac_key = hmac_key(key, first[:16])
    if not matches(first, key):
        raise ValueError('密钥或加密参数不匹配')
    if source.stat().st_size % PAGE:
        raise ValueError('源数据库大小不完整')
    frames, count = wal_frames(Path(str(source) + '-wal'))
    created = False
    try:
        with source.open('rb') as file, destination.open('xb') as out:
            created = True
            try:
                destination.chmod(0o600)
            except OSError:
                pass
            number = 0
            while page := file.read(PAGE):
                number += 1
                out.write(decrypt_page(page, number, key, mac_key))
            if frames:
                with Path(str(source) + '-wal').open('rb') as wal:
                    for offset, number in frames:
                        wal.seek(offset)
                        plain = decrypt_page(wal.read(PAGE), number, key, mac_key)
                        out.seek((number - 1) * PAGE)
                        out.write(plain)
                out.truncate(count * PAGE)
            # WAL is now replayed into this standalone, private output only.
            out.seek(18)
            out.write(b'\x01\x01')
        db = sqlite3.connect(destination.resolve().as_uri() + '?mode=ro&immutable=1', uri=True)
        try:
            result = db.execute('PRAGMA quick_check').fetchall()
            if result != [('ok',)]:
                raise ValueError('解密后的数据库完整性检查未通过')
        finally:
            db.close()
    except Exception:
        if created:
            destination.unlink(missing_ok=True)
        raise
    return {'wal_frames': len(frames), 'pages': destination.stat().st_size // PAGE}