"""Формат зашифрованного файла базы (.fmdb).

Вся база — один JSON-документ. На диск он попадает так:

    json → utf-8 → gzip → AES-256-GCM → файл

Ключ выводится из пароля через scrypt. Заголовок файла (37 байт) передаётся
в GCM как AAD, поэтому подменить параметры KDF или соль незаметно нельзя —
расшифровка упадёт с InvalidTag так же, как при неверном пароле.

    смещение  размер  содержимое
    0         4       магия b"FMDB"
    4         1       версия формата
    5         1       id функции выведения ключа (1 = scrypt)
    6         3       параметры scrypt: log2(n), r, p
    9         16      salt
    25        12      nonce
    37        …       ciphertext + тег GCM (16 байт)
"""

from __future__ import annotations

import gzip
import os
import struct

from cryptography.exceptions import InvalidTag
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt

MAGIC = b"FMDB"
FORMAT_VERSION = 1
KDF_SCRYPT = 1

SALT_SIZE = 16
NONCE_SIZE = 12
KEY_SIZE = 32
HEADER_SIZE = 37

# n=2**15 — примерно 100 мс на вывод ключа. Достаточно, чтобы перебор паролей
# был дорогим, и незаметно при открытии базы.
SCRYPT_LOG_N = 15
SCRYPT_R = 8
SCRYPT_P = 1

_HEADER_STRUCT = struct.Struct(f"!4sBBBBB{SALT_SIZE}s{NONCE_SIZE}s")
assert _HEADER_STRUCT.size == HEADER_SIZE


class VaultError(Exception):
    """Базовая ошибка работы с файлом базы."""


class WrongPassword(VaultError):
    """Пароль неверный либо файл повреждён/подделан.

    Намеренно не различаем эти случаи: по тексту ошибки не должно быть видно,
    угадана ли часть пароля.
    """


class UnsupportedFormat(VaultError):
    """Файл создан более новой версией программы либо это вообще не база."""


class Header:
    """Разобранный заголовок файла."""

    __slots__ = ("version", "kdf_id", "log_n", "r", "p", "salt", "nonce", "raw")

    def __init__(
        self,
        version: int,
        kdf_id: int,
        log_n: int,
        r: int,
        p: int,
        salt: bytes,
        nonce: bytes,
    ) -> None:
        self.version = version
        self.kdf_id = kdf_id
        self.log_n = log_n
        self.r = r
        self.p = p
        self.salt = salt
        self.nonce = nonce
        self.raw = _HEADER_STRUCT.pack(
            MAGIC, version, kdf_id, log_n, r, p, salt, nonce
        )

    @classmethod
    def new(cls) -> "Header":
        """Свежий заголовок со случайными salt и nonce."""
        return cls(
            version=FORMAT_VERSION,
            kdf_id=KDF_SCRYPT,
            log_n=SCRYPT_LOG_N,
            r=SCRYPT_R,
            p=SCRYPT_P,
            salt=os.urandom(SALT_SIZE),
            nonce=os.urandom(NONCE_SIZE),
        )

    @classmethod
    def parse(cls, blob: bytes) -> "Header":
        if len(blob) < HEADER_SIZE:
            raise UnsupportedFormat("Файл слишком короткий, это не база.")

        magic, version, kdf_id, log_n, r, p, salt, nonce = _HEADER_STRUCT.unpack(
            blob[:HEADER_SIZE]
        )
        if magic != MAGIC:
            raise UnsupportedFormat("Это не файл базы (неверная сигнатура).")
        if version > FORMAT_VERSION:
            raise UnsupportedFormat(
                f"Файл создан более новой версией программы (формат {version}). "
                "Обнови приложение."
            )
        if kdf_id != KDF_SCRYPT:
            raise UnsupportedFormat(f"Неизвестный способ выведения ключа: {kdf_id}.")
        # Верхняя граница защищает от файла, который просит гигабайты памяти.
        if not 10 <= log_n <= 22:
            raise UnsupportedFormat(f"Недопустимый параметр scrypt: log2(n)={log_n}.")
        if not 1 <= r <= 32 or not 1 <= p <= 16:
            raise UnsupportedFormat("Недопустимые параметры scrypt.")

        return cls(version, kdf_id, log_n, r, p, salt, nonce)


def derive_key(password: str, header: Header) -> bytes:
    """Вывести ключ шифрования из пароля по параметрам из заголовка."""
    kdf = Scrypt(
        salt=header.salt,
        length=KEY_SIZE,
        n=2**header.log_n,
        r=header.r,
        p=header.p,
    )
    return kdf.derive(password.encode("utf-8"))


def encrypt(plaintext: bytes, password: str) -> bytes:
    """Сжать и зашифровать открытые данные, вернуть содержимое файла целиком.

    Каждый вызов берёт новые salt и nonce, поэтому одни и те же данные дают
    разные байты. Проверять «изменилось ли что-то» по этому результату нельзя —
    для этого storage хранит хеш открытого JSON.
    """
    header = Header.new()
    key = derive_key(password, header)
    packed = gzip.compress(plaintext, compresslevel=9, mtime=0)
    ciphertext = AESGCM(key).encrypt(header.nonce, packed, header.raw)
    return header.raw + ciphertext


def decrypt(blob: bytes, password: str) -> bytes:
    """Расшифровать содержимое файла, вернуть открытые данные."""
    header = Header.parse(blob)
    key = derive_key(password, header)
    try:
        packed = AESGCM(key).decrypt(header.nonce, blob[HEADER_SIZE:], header.raw)
    except InvalidTag as exc:
        raise WrongPassword("Неверный пароль или файл повреждён.") from exc

    try:
        return gzip.decompress(packed)
    except (OSError, EOFError) as exc:
        # Сюда попасть почти невозможно: GCM уже подтвердил целостность.
        raise VaultError("Данные расшифровались, но распаковать их не удалось.") from exc