document-management/backend/tests/dms/test_envelope.py

"""契约1.21 信封加密（KEK/DEK，AES-256-GCM）单元测试。

覆盖：加解密往返、DMSE 识别、明文兼容、promote 门控（original/ 前缀
与 KEK 配置双条件）、篡改检测、缺 KEK 拒绝读取、KEK 配置解析校验。
"""

from __future__ import annotations

import base64
import os
from pathlib import Path

import pytest
from flask import Flask

from dms.common.errors import InternalError
from dms.storage import envelope


@pytest.fixture()
def kek() -> bytes:
    return os.urandom(32)


@pytest.fixture()
def crypto_app(kek):
    app = Flask("dms-envelope-test")
    app.config["DMS_STORAGE_KEK_BYTES"] = kek
    return app


@pytest.fixture()
def plain_app():
    app = Flask("dms-envelope-plain-test")
    app.config["DMS_STORAGE_KEK_BYTES"] = None
    return app


class TestLoadKek:
    def test_empty_returns_none(self):
        assert envelope.load_kek("") is None
        assert envelope.load_kek("   ") is None

    def test_valid_base64_32_bytes(self):
        raw = os.urandom(32)
        assert envelope.load_kek(base64.b64encode(raw).decode()) == raw

    def test_invalid_base64_rejected(self):
        with pytest.raises(ValueError, match="DMS_STORAGE_KEK"):
            envelope.load_kek("!!not-base64!!")

    def test_wrong_length_rejected(self):
        with pytest.raises(ValueError, match="DMS_STORAGE_KEK"):
            envelope.load_kek(base64.b64encode(os.urandom(16)).decode())


class TestEncryptDecryptRoundTrip:
    def test_round_trip(self, crypto_app, kek, tmp_path: Path):
        plaintext = os.urandom(4096)
        source = tmp_path / "plain.bin"
        source.write_bytes(plaintext)
        target = tmp_path / "encrypted.bin"

        envelope.encrypt_file(kek, source, target)

        payload = target.read_bytes()
        assert payload[:4] == envelope.ENVELOPE_MAGIC
        assert payload[4] == envelope.KEK_VERSION
        assert envelope.decrypt_bytes(kek, payload) == plaintext

    def test_each_file_uses_distinct_dek(self, crypto_app, kek, tmp_path: Path):
        source = tmp_path / "plain.bin"
        source.write_bytes(b"same content")
        first = tmp_path / "a.bin"
        second = tmp_path / "b.bin"
        envelope.encrypt_file(kek, source, first)
        envelope.encrypt_file(kek, source, second)
        # DEK/nonce 随机，同明文两次加密产物必须不同
        assert first.read_bytes() != second.read_bytes()

    def test_wrong_kek_rejected(self, crypto_app, kek, tmp_path: Path):
        source = tmp_path / "plain.bin"
        source.write_bytes(b"secret")
        target = tmp_path / "encrypted.bin"
        envelope.encrypt_file(kek, source, target)
        with pytest.raises(InternalError, match="解密失败"):
            envelope.decrypt_bytes(os.urandom(32), target.read_bytes())

    def test_tampered_ciphertext_rejected(self, crypto_app, kek, tmp_path: Path):
        source = tmp_path / "plain.bin"
        source.write_bytes(b"integrity check")
        target = tmp_path / "encrypted.bin"
        envelope.encrypt_file(kek, source, target)
        payload = bytearray(target.read_bytes())
        payload[-1] ^= 0xFF
        with pytest.raises(InternalError, match="解密失败"):
            envelope.decrypt_bytes(kek, bytes(payload))

    def test_truncated_payload_rejected(self, crypto_app, kek):
        with pytest.raises(InternalError, match="解密失败"):
            envelope.decrypt_bytes(kek, b"DMSE")
        with pytest.raises(InternalError, match="解密失败"):
            envelope.decrypt_bytes(kek, b"NOPE" + b"\x00" * 64)


class TestIsEnvelope:
    def test_recognizes_envelope(self, crypto_app, kek, tmp_path: Path):
        source = tmp_path / "plain.bin"
        source.write_bytes(b"data")
        target = tmp_path / "encrypted.bin"
        envelope.encrypt_file(kek, source, target)
        assert envelope.is_envelope(target) is True

    def test_plain_file_not_envelope(self, tmp_path: Path):
        path = tmp_path / "plain.pdf"
        path.write_bytes(b"%PDF-1.7 body")
        assert envelope.is_envelope(path) is False

    def test_missing_file_not_envelope(self, tmp_path: Path):
        assert envelope.is_envelope(tmp_path / "missing.bin") is False


class TestReadPlaintextBytes:
    def test_plain_passthrough(self, plain_app, tmp_path: Path):
        path = tmp_path / "plain.pdf"
        content = b"%PDF-1.7 plain content"
        path.write_bytes(content)
        with plain_app.app_context():
            assert envelope.read_plaintext_bytes(path) == content

    def test_envelope_decrypted(self, crypto_app, kek, tmp_path: Path):
        plaintext = b"confidential bytes"
        source = tmp_path / "plain.bin"
        source.write_bytes(plaintext)
        target = tmp_path / "encrypted.bin"
        envelope.encrypt_file(kek, source, target)
        with crypto_app.app_context():
            assert envelope.read_plaintext_bytes(target) == plaintext

    def test_envelope_without_kek_rejected(self, plain_app, kek, tmp_path: Path):
        source = tmp_path / "plain.bin"
        source.write_bytes(b"encrypted at rest")
        target = tmp_path / "encrypted.bin"
        envelope.encrypt_file(kek, source, target)
        with plain_app.app_context():
            with pytest.raises(InternalError, match="未配置"):
                envelope.read_plaintext_bytes(target)


class TestPlaintextFileFor:
    def test_plain_returns_same_path(self, plain_app, tmp_path: Path):
        path = tmp_path / "plain.pdf"
        path.write_bytes(b"%PDF-1.7 body")
        with plain_app.app_context():
            with envelope.plaintext_file_for(path) as yielded:
                assert yielded == path
        assert path.exists()

    def test_envelope_yields_temp_and_cleans_up(
        self, crypto_app, kek, tmp_path: Path
    ):
        plaintext = b"decrypt me"
        source = tmp_path / "plain.bin"
        source.write_bytes(plaintext)
        target = tmp_path / "encrypted.bin"
        envelope.encrypt_file(kek, source, target)
        with crypto_app.app_context():
            with envelope.plaintext_file_for(target) as yielded:
                assert yielded != target
                assert yielded.read_bytes() == plaintext
            assert not yielded.exists()
        # 密文本体不受临时明文清理影响
        assert target.exists()


class TestPromoteEncrypt:
    def test_original_prefix_encrypted(self, crypto_app, tmp_path: Path):
        temporary = tmp_path / "staged.upload"
        temporary.write_bytes(b"new upload")
        final = tmp_path / "final.bin"
        with crypto_app.app_context():
            assert envelope.promote_encrypt(
                temporary, final, "original/final.bin"
            ) is True
        assert final.read_bytes()[:4] == envelope.ENVELOPE_MAGIC
        assert not temporary.exists()

    def test_non_original_prefix_skipped(self, crypto_app, tmp_path: Path):
        temporary = tmp_path / "staged.upload"
        temporary.write_bytes(b"voucher pdf bytes")
        final = tmp_path / "final.pdf"
        with crypto_app.app_context():
            assert envelope.promote_encrypt(
                temporary, final, "voucher/final.pdf"
            ) is False
        assert not final.exists()
        assert temporary.exists()

    def test_no_kek_skipped(self, plain_app, tmp_path: Path):
        temporary = tmp_path / "staged.upload"
        temporary.write_bytes(b"plain stays plain")
        final = tmp_path / "final.bin"
        with plain_app.app_context():
            assert envelope.promote_encrypt(
                temporary, final, "original/final.bin"
            ) is False
        assert not final.exists()
        assert temporary.exists()