PROJECT-CONTORL/backend/tests/test_encryption.py

"""
Tests for the file encryption functionality.

Tests cover:
- Encryption service (key generation, encrypt/decrypt)
- Encryption key management API
- Attachment upload with encryption
- Attachment download with decryption
"""
import pytest
import base64
import secrets
from io import BytesIO
from unittest.mock import patch, MagicMock

from app.services.encryption_service import (
    EncryptionService,
    encryption_service,
    MasterKeyNotConfiguredError,
    DecryptionError,
    KEY_SIZE,
    NONCE_SIZE,
)


class TestEncryptionService:
    """Tests for the encryption service."""

    @pytest.fixture
    def mock_master_key(self):
        """Generate a valid test master key."""
        return base64.urlsafe_b64encode(secrets.token_bytes(32)).decode()

    @pytest.fixture
    def service_with_key(self, mock_master_key):
        """Create an encryption service with a mock master key."""
        with patch('app.services.encryption_service.settings') as mock_settings:
            mock_settings.ENCRYPTION_MASTER_KEY = mock_master_key
            service = EncryptionService()
            yield service

    def test_generate_key(self, service_with_key):
        """Test that generate_key produces a 32-byte key."""
        key = service_with_key.generate_key()
        assert len(key) == KEY_SIZE
        assert isinstance(key, bytes)

    def test_generate_key_uniqueness(self, service_with_key):
        """Test that each generated key is unique."""
        keys = [service_with_key.generate_key() for _ in range(10)]
        unique_keys = set(keys)
        assert len(unique_keys) == 10

    def test_encrypt_decrypt_key(self, service_with_key):
        """Test encryption and decryption of a file encryption key."""
        # Generate a key to encrypt
        original_key = service_with_key.generate_key()

        # Encrypt the key
        encrypted_key = service_with_key.encrypt_key(original_key)
        assert isinstance(encrypted_key, str)
        assert encrypted_key != base64.urlsafe_b64encode(original_key).decode()

        # Decrypt the key
        decrypted_key = service_with_key.decrypt_key(encrypted_key)
        assert decrypted_key == original_key

    def test_encrypt_decrypt_file(self, service_with_key):
        """Test file encryption and decryption."""
        # Create test file content
        original_content = b"This is a test file content for encryption."
        file_obj = BytesIO(original_content)

        # Generate encryption key
        key = service_with_key.generate_key()

        # Encrypt
        encrypted_content = service_with_key.encrypt_file(file_obj, key)
        assert encrypted_content != original_content
        assert len(encrypted_content) > len(original_content)  # Due to nonce and tag

        # Decrypt
        encrypted_file = BytesIO(encrypted_content)
        decrypted_content = service_with_key.decrypt_file(encrypted_file, key)
        assert decrypted_content == original_content

    def test_encrypt_decrypt_bytes(self, service_with_key):
        """Test bytes encryption and decryption convenience methods."""
        original_data = b"Test data for encryption"
        key = service_with_key.generate_key()

        # Encrypt
        encrypted_data = service_with_key.encrypt_bytes(original_data, key)
        assert encrypted_data != original_data

        # Decrypt
        decrypted_data = service_with_key.decrypt_bytes(encrypted_data, key)
        assert decrypted_data == original_data

    def test_encrypt_large_file(self, service_with_key):
        """Test encryption of a larger file (1MB)."""
        # Create 1MB of random data
        original_content = secrets.token_bytes(1024 * 1024)
        file_obj = BytesIO(original_content)
        key = service_with_key.generate_key()

        # Encrypt
        encrypted_content = service_with_key.encrypt_file(file_obj, key)

        # Decrypt
        encrypted_file = BytesIO(encrypted_content)
        decrypted_content = service_with_key.decrypt_file(encrypted_file, key)

        assert decrypted_content == original_content

    def test_decrypt_with_wrong_key(self, service_with_key):
        """Test that decryption fails with wrong key."""
        original_content = b"Secret content"
        file_obj = BytesIO(original_content)

        key1 = service_with_key.generate_key()
        key2 = service_with_key.generate_key()

        # Encrypt with key1
        encrypted_content = service_with_key.encrypt_file(file_obj, key1)

        # Try to decrypt with key2
        encrypted_file = BytesIO(encrypted_content)
        with pytest.raises(DecryptionError):
            service_with_key.decrypt_file(encrypted_file, key2)

    def test_decrypt_corrupted_data(self, service_with_key):
        """Test that decryption fails with corrupted data."""
        original_content = b"Secret content"
        file_obj = BytesIO(original_content)
        key = service_with_key.generate_key()

        # Encrypt
        encrypted_content = service_with_key.encrypt_file(file_obj, key)

        # Corrupt the encrypted data
        corrupted = bytearray(encrypted_content)
        corrupted[20] ^= 0xFF  # Flip some bits
        corrupted_content = bytes(corrupted)

        # Try to decrypt
        encrypted_file = BytesIO(corrupted_content)
        with pytest.raises(DecryptionError):
            service_with_key.decrypt_file(encrypted_file, key)

    def test_is_encryption_available_with_key(self, mock_master_key):
        """Test is_encryption_available returns True when key is configured."""
        with patch('app.services.encryption_service.settings') as mock_settings:
            mock_settings.ENCRYPTION_MASTER_KEY = mock_master_key
            service = EncryptionService()
            assert service.is_encryption_available() is True

    def test_is_encryption_available_without_key(self):
        """Test is_encryption_available returns False when key is not configured."""
        with patch('app.services.encryption_service.settings') as mock_settings:
            mock_settings.ENCRYPTION_MASTER_KEY = None
            service = EncryptionService()
            assert service.is_encryption_available() is False

    def test_master_key_not_configured_error(self):
        """Test that operations fail when master key is not configured."""
        with patch('app.services.encryption_service.settings') as mock_settings:
            mock_settings.ENCRYPTION_MASTER_KEY = None
            service = EncryptionService()

            key = secrets.token_bytes(32)
            with pytest.raises(MasterKeyNotConfiguredError):
                service.encrypt_key(key)

    def test_encrypted_key_format(self, service_with_key):
        """Test that encrypted key is valid base64."""
        key = service_with_key.generate_key()
        encrypted_key = service_with_key.encrypt_key(key)

        # Should be valid base64
        decoded = base64.urlsafe_b64decode(encrypted_key)
        # Should contain nonce + ciphertext + tag
        assert len(decoded) >= NONCE_SIZE + KEY_SIZE + 16  # 16 = GCM tag size


class TestEncryptionServiceStreaming:
    """Tests for streaming encryption (for large files)."""

    @pytest.fixture
    def mock_master_key(self):
        """Generate a valid test master key."""
        return base64.urlsafe_b64encode(secrets.token_bytes(32)).decode()

    @pytest.fixture
    def service_with_key(self, mock_master_key):
        """Create an encryption service with a mock master key."""
        with patch('app.services.encryption_service.settings') as mock_settings:
            mock_settings.ENCRYPTION_MASTER_KEY = mock_master_key
            service = EncryptionService()
            yield service

    def test_streaming_encrypt_decrypt(self, service_with_key):
        """Test streaming encryption and decryption."""
        # Create test content
        original_content = b"Test content for streaming encryption. " * 1000
        file_obj = BytesIO(original_content)
        key = service_with_key.generate_key()

        # Encrypt using streaming
        encrypted_chunks = list(service_with_key.encrypt_file_streaming(file_obj, key))
        encrypted_content = b''.join(encrypted_chunks)

        # Decrypt using streaming
        encrypted_file = BytesIO(encrypted_content)
        decrypted_chunks = list(service_with_key.decrypt_file_streaming(encrypted_file, key))
        decrypted_content = b''.join(decrypted_chunks)

        assert decrypted_content == original_content


class TestEncryptionKeyValidation:
    """Tests for encryption key validation in config."""

    def test_valid_master_key(self):
        """Test that a valid master key passes validation."""
        from app.core.config import Settings

        valid_key = base64.urlsafe_b64encode(secrets.token_bytes(32)).decode()

        # This should not raise
        with patch.dict('os.environ', {
            'JWT_SECRET_KEY': 'test-secret-key-that-is-valid',
            'ENCRYPTION_MASTER_KEY': valid_key
        }):
            settings = Settings()
            assert settings.ENCRYPTION_MASTER_KEY == valid_key

    def test_invalid_master_key_length(self):
        """Test that an invalid length master key fails validation."""
        from app.core.config import Settings

        # 16 bytes instead of 32
        invalid_key = base64.urlsafe_b64encode(secrets.token_bytes(16)).decode()

        with patch.dict('os.environ', {
            'JWT_SECRET_KEY': 'test-secret-key-that-is-valid',
            'ENCRYPTION_MASTER_KEY': invalid_key
        }):
            with pytest.raises(ValueError, match="must be a base64-encoded 32-byte key"):
                Settings()

    def test_invalid_master_key_format(self):
        """Test that an invalid format master key fails validation."""
        from app.core.config import Settings
        from pydantic import ValidationError

        invalid_key = "not-valid-base64!@#$"

        with patch.dict('os.environ', {
            'JWT_SECRET_KEY': 'test-secret-key-that-is-valid',
            'ENCRYPTION_MASTER_KEY': invalid_key
        }):
            with pytest.raises(ValidationError, match="ENCRYPTION_MASTER_KEY"):
                Settings()

    def test_empty_master_key_allowed(self):
        """Test that empty master key is allowed (encryption disabled)."""
        from app.core.config import Settings

        with patch.dict('os.environ', {
            'JWT_SECRET_KEY': 'test-secret-key-that-is-valid',
            'ENCRYPTION_MASTER_KEY': ''
        }):
            settings = Settings()
            assert settings.ENCRYPTION_MASTER_KEY is None