document-management/backend/dms/storage/envelope.py

"""应用层信封加密（KEK/DEK，AES-256-GCM）。

契约 1.21：配置 DMS_STORAGE_KEK（base64 编码的 32 字节密钥）后，
original/ 新落盘文件按以下格式加密存储：
  DMSE + kekVersion(1B) + 包裹DEK长度(4B, 大端) + KEK包裹的DEK + 密文
其中：
- "KEK包裹的DEK" = 包裹nonce(12B) + AES-256-GCM(KEK) 加密的 DEK 密文（含认证标签）
- "密文" = 数据nonce(12B) + AES-256-GCM(DEK) 加密的文件明文（含认证标签）
- DEK 每文件随机 32 字节，nonce 均 12 字节
读路径按 DMSE 头自动识别密文/明文实现明文兼容；
未配置 KEK 时维持明文落盘，存量文件不回迁加密。
"""

from __future__ import annotations

import base64
import os
import struct
import tempfile
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator

from cryptography.exceptions import InvalidTag
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from flask import current_app

from dms.common.errors import InternalError

ENVELOPE_MAGIC = b"DMSE"
KEK_VERSION = 1
_KEK_LENGTH = 32
_DEK_LENGTH = 32
_NONCE_LENGTH = 12
_GCM_TAG_LENGTH = 16
# 包裹 DEK 长度防御性上限（正常为 12+32+16=60 字节）
_WRAPPED_DEK_MAX = 1024


def load_kek(raw_value: str) -> bytes | None:
    """解析并校验 DMS_STORAGE_KEK：base64 编码的 32 字节密钥；空值返回 None。"""
    value = raw_value.strip()
    if not value:
        return None
    try:
        kek = base64.b64decode(value, validate=True)
    except ValueError as exc:
        raise ValueError("DMS_STORAGE_KEK必须是base64编码的32字节密钥") from exc
    if len(kek) != _KEK_LENGTH:
        raise ValueError("DMS_STORAGE_KEK必须是base64编码的32字节密钥")
    return kek


def _current_kek() -> bytes | None:
    return current_app.config.get("DMS_STORAGE_KEK_BYTES")


def _require_kek() -> bytes:
    kek = _current_kek()
    if kek is None:
        raise InternalError("存储加密密钥未配置，无法读取加密文件")
    return kek


def is_envelope(path: Path) -> bool:
    """按 DMSE 魔数识别信封加密文件。"""
    try:
        with path.open("rb") as stream:
            return stream.read(len(ENVELOPE_MAGIC)) == ENVELOPE_MAGIC
    except OSError:
        return False


def encrypt_file(kek: bytes, source: Path, target: Path) -> None:
    """将明文文件加密为信封格式落盘。"""
    plaintext = source.read_bytes()
    dek = os.urandom(_DEK_LENGTH)
    wrap_nonce = os.urandom(_NONCE_LENGTH)
    data_nonce = os.urandom(_NONCE_LENGTH)
    wrapped_dek = wrap_nonce + AESGCM(kek).encrypt(wrap_nonce, dek, None)
    ciphertext = data_nonce + AESGCM(dek).encrypt(data_nonce, plaintext, None)
    header = (
        ENVELOPE_MAGIC
        + bytes([KEK_VERSION])
        + struct.pack(">I", len(wrapped_dek))
        + wrapped_dek
    )
    target.write_bytes(header + ciphertext)


def decrypt_bytes(kek: bytes, payload: bytes) -> bytes:
    """解迷信封格式字节串，返回明文；格式损坏或认证失败统一报内部错误。"""
    try:
        if len(payload) < len(ENVELOPE_MAGIC) + 1 + 4:
            raise ValueError("envelope too short")
        if payload[: len(ENVELOPE_MAGIC)] != ENVELOPE_MAGIC:
            raise ValueError("bad magic")
        if payload[4] != KEK_VERSION:
            raise ValueError("unsupported kek version")
        (wrapped_length,) = struct.unpack(">I", payload[5:9])
        if not _NONCE_LENGTH < wrapped_length <= _WRAPPED_DEK_MAX:
            raise ValueError("bad wrapped dek length")
        offset = 9
        wrapped = payload[offset : offset + wrapped_length]
        offset += wrapped_length
        ciphertext = payload[offset:]
        if len(ciphertext) < _NONCE_LENGTH + _GCM_TAG_LENGTH:
            raise ValueError("ciphertext too short")
        wrap_nonce, wrapped_dek = (
            wrapped[:_NONCE_LENGTH],
            wrapped[_NONCE_LENGTH:],
        )
        dek = AESGCM(kek).decrypt(wrap_nonce, wrapped_dek, None)
        data_nonce, data = (
            ciphertext[:_NONCE_LENGTH],
            ciphertext[_NONCE_LENGTH:],
        )
        return AESGCM(dek).decrypt(data_nonce, data, None)
    except (ValueError, InvalidTag) as exc:
        raise InternalError("文档文件解密失败或已损坏") from exc


def read_plaintext_bytes(path: Path) -> bytes:
    """读取文件明文：DMSE 头密文自动解密，否则按明文原样返回。"""
    payload = path.read_bytes()
    if not payload.startswith(ENVELOPE_MAGIC):
        return payload
    return decrypt_bytes(_require_kek(), payload)


@contextmanager
def plaintext_file_for(path: Path) -> Iterator[Path]:
    """提供明文文件路径：明文原样返回；密文解密至临时文件，退出后清理。"""
    if not is_envelope(path):
        yield path
        return
    plaintext = read_plaintext_bytes(path)
    handle = tempfile.NamedTemporaryFile(delete=False)
    try:
        handle.write(plaintext)
        handle.close()
        yield Path(handle.name)
    finally:
        Path(handle.name).unlink(missing_ok=True)


def promote_encrypt(temporary: Path, final: Path, relative_path: str) -> bool:
    """promote 门控：original/ 且已配置 KEK 时加密落盘并清理暂存，返回 True。"""
    kek = _current_kek()
    if kek is None or not relative_path.startswith("original/"):
        return False
    encrypt_file(kek, temporary, final)
    temporary.unlink(missing_ok=True)
    return True