317 lines
10 KiB
Python
317 lines
10 KiB
Python
"""Unit tests for encrypted storage functionality."""
|
|
|
|
import json
|
|
import tempfile
|
|
from pathlib import Path
|
|
from unittest.mock import patch, mock_open
|
|
import pytest
|
|
|
|
from src.security.encrypted_storage import (
|
|
EncryptedStorage,
|
|
encrypt_data,
|
|
decrypt_data,
|
|
generate_encryption_key,
|
|
)
|
|
|
|
|
|
class TestEncryptedStorage:
|
|
"""Test cases for EncryptedStorage class."""
|
|
|
|
def setup_method(self):
|
|
"""Set up test fixtures."""
|
|
self.temp_dir = tempfile.mkdtemp()
|
|
self.storage_path = Path(self.temp_dir) / "encrypted_data"
|
|
self.key_path = Path(self.temp_dir) / "key.bin"
|
|
|
|
def teardown_method(self):
|
|
"""Clean up test fixtures."""
|
|
import shutil
|
|
if self.temp_dir and Path(self.temp_dir).exists():
|
|
shutil.rmtree(self.temp_dir)
|
|
|
|
def test_init_creates_storage_directory(self):
|
|
"""Test that EncryptedStorage creates storage directory if it doesn't exist."""
|
|
storage_dir = Path(self.temp_dir) / "new_storage"
|
|
storage_path = storage_dir / "data"
|
|
|
|
# Directory shouldn't exist initially
|
|
assert not storage_dir.exists()
|
|
|
|
# Create EncryptedStorage instance
|
|
storage = EncryptedStorage(storage_path, self.key_path)
|
|
|
|
# Directory should be created
|
|
assert storage_dir.exists()
|
|
assert storage_dir.is_dir()
|
|
|
|
def test_init_generates_new_key_if_not_exists(self):
|
|
"""Test that EncryptedStorage generates a new encryption key if it doesn't exist."""
|
|
# Key file shouldn't exist initially
|
|
assert not self.key_path.exists()
|
|
|
|
# Create EncryptedStorage instance
|
|
storage = EncryptedStorage(self.storage_path, self.key_path)
|
|
|
|
# Key file should be created
|
|
assert self.key_path.exists()
|
|
assert self.key_path.is_file()
|
|
|
|
def test_store_and_retrieve_string_data(self):
|
|
"""Test storing and retrieving string data."""
|
|
storage = EncryptedStorage(self.storage_path, self.key_path)
|
|
|
|
# Store string data
|
|
test_data = "sensitive string data"
|
|
key = "test_key"
|
|
|
|
success = storage.store(key, test_data)
|
|
assert success is True
|
|
|
|
# Retrieve data
|
|
retrieved_data = storage.retrieve(key)
|
|
assert retrieved_data == test_data
|
|
|
|
def test_store_and_retrieve_dict_data(self):
|
|
"""Test storing and retrieving dictionary data."""
|
|
storage = EncryptedStorage(self.storage_path, self.key_path)
|
|
|
|
# Store dictionary data
|
|
test_data = {"user_id": 123, "api_key": "secret_key", "settings": {"theme": "dark"}}
|
|
key = "user_config"
|
|
|
|
success = storage.store(key, test_data)
|
|
assert success is True
|
|
|
|
# Retrieve data
|
|
retrieved_data = storage.retrieve(key)
|
|
assert retrieved_data == test_data
|
|
|
|
def test_store_and_retrieve_list_data(self):
|
|
"""Test storing and retrieving list data."""
|
|
storage = EncryptedStorage(self.storage_path, self.key_path)
|
|
|
|
# Store list data
|
|
test_data = ["item1", "item2", {"nested": "data"}, 123]
|
|
key = "list_data"
|
|
|
|
success = storage.store(key, test_data)
|
|
assert success is True
|
|
|
|
# Retrieve data
|
|
retrieved_data = storage.retrieve(key)
|
|
assert retrieved_data == test_data
|
|
|
|
def test_store_and_retrieve_binary_data(self):
|
|
"""Test storing and retrieving binary data."""
|
|
storage = EncryptedStorage(self.storage_path, self.key_path)
|
|
|
|
# Store binary data
|
|
test_data = b"binary sensitive data"
|
|
key = "binary_key"
|
|
|
|
success = storage.store(key, test_data)
|
|
assert success is True
|
|
|
|
# Retrieve data
|
|
retrieved_data = storage.retrieve(key)
|
|
assert retrieved_data == test_data
|
|
|
|
def test_retrieve_nonexistent_key(self):
|
|
"""Test retrieving data for a key that doesn't exist."""
|
|
storage = EncryptedStorage(self.storage_path, self.key_path)
|
|
|
|
# Try to retrieve non-existent key
|
|
result = storage.retrieve("nonexistent_key")
|
|
assert result is None
|
|
|
|
def test_overwrite_existing_key(self):
|
|
"""Test overwriting data for an existing key."""
|
|
storage = EncryptedStorage(self.storage_path, self.key_path)
|
|
|
|
# Store initial data
|
|
key = "test_key"
|
|
initial_data = "initial data"
|
|
storage.store(key, initial_data)
|
|
|
|
# Overwrite with new data
|
|
new_data = "new data"
|
|
success = storage.store(key, new_data)
|
|
assert success is True
|
|
|
|
# Retrieve should return new data
|
|
retrieved_data = storage.retrieve(key)
|
|
assert retrieved_data == new_data
|
|
|
|
def test_delete_existing_key(self):
|
|
"""Test deleting data for an existing key."""
|
|
storage = EncryptedStorage(self.storage_path, self.key_path)
|
|
|
|
# Store data
|
|
key = "test_key"
|
|
test_data = "test data"
|
|
storage.store(key, test_data)
|
|
|
|
# Verify data exists
|
|
assert storage.retrieve(key) == test_data
|
|
|
|
# Delete data
|
|
success = storage.delete(key)
|
|
assert success is True
|
|
|
|
# Verify data is gone
|
|
assert storage.retrieve(key) is None
|
|
|
|
def test_delete_nonexistent_key(self):
|
|
"""Test deleting data for a key that doesn't exist."""
|
|
storage = EncryptedStorage(self.storage_path, self.key_path)
|
|
|
|
# Try to delete non-existent key
|
|
success = storage.delete("nonexistent_key")
|
|
assert success is False
|
|
|
|
def test_list_keys(self):
|
|
"""Test listing all stored keys."""
|
|
storage = EncryptedStorage(self.storage_path, self.key_path)
|
|
|
|
# Store multiple items
|
|
storage.store("key1", "data1")
|
|
storage.store("key2", "data2")
|
|
storage.store("key3", "data3")
|
|
|
|
# List keys
|
|
keys = storage.list_keys()
|
|
assert "key1" in keys
|
|
assert "key2" in keys
|
|
assert "key3" in keys
|
|
assert len(keys) == 3
|
|
|
|
def test_clear_all_data(self):
|
|
"""Test clearing all stored data."""
|
|
storage = EncryptedStorage(self.storage_path, self.key_path)
|
|
|
|
# Store multiple items
|
|
storage.store("key1", "data1")
|
|
storage.store("key2", "data2")
|
|
|
|
# Verify data exists
|
|
assert len(storage.list_keys()) == 2
|
|
|
|
# Clear all data
|
|
success = storage.clear()
|
|
assert success is True
|
|
|
|
# Verify all data is gone
|
|
assert len(storage.list_keys()) == 0
|
|
assert storage.retrieve("key1") is None
|
|
assert storage.retrieve("key2") is None
|
|
|
|
def test_handle_encryption_errors(self):
|
|
"""Test handling encryption errors gracefully."""
|
|
storage = EncryptedStorage(self.storage_path, self.key_path)
|
|
|
|
# Mock encryption to fail
|
|
with patch.object(storage.fernet, 'encrypt', side_effect=Exception("Encryption failed")):
|
|
success = storage.store("test_key", "test_data")
|
|
assert success is False
|
|
|
|
def test_handle_decryption_errors(self):
|
|
"""Test handling decryption errors gracefully."""
|
|
storage = EncryptedStorage(self.storage_path, self.key_path)
|
|
|
|
# Store valid data first
|
|
storage.store("test_key", "test_data")
|
|
|
|
# Corrupt the stored data
|
|
storage_file = self.storage_path / "test_key.enc"
|
|
with open(storage_file, "wb") as f:
|
|
f.write(b"corrupted_data")
|
|
|
|
# Try to retrieve corrupted data
|
|
result = storage.retrieve("test_key")
|
|
assert result is None
|
|
|
|
def test_file_permissions(self):
|
|
"""Test that files have correct permissions."""
|
|
storage = EncryptedStorage(self.storage_path, self.key_path)
|
|
|
|
# Store data
|
|
storage.store("test_key", "test_data")
|
|
|
|
# Check key file permissions
|
|
key_stat = self.key_path.stat()
|
|
assert oct(key_stat.st_mode)[-3:] == "600"
|
|
|
|
# Check storage file permissions
|
|
storage_file = self.storage_path / "test_key.enc"
|
|
storage_stat = storage_file.stat()
|
|
assert oct(storage_stat.st_mode)[-3:] == "600"
|
|
|
|
|
|
class TestEncryptionUtilities:
|
|
"""Test cases for encryption utility functions."""
|
|
|
|
def test_generate_encryption_key(self):
|
|
"""Test generating encryption key."""
|
|
key = generate_encryption_key()
|
|
assert len(key) == 44 # Fernet key length
|
|
assert isinstance(key, bytes)
|
|
|
|
def test_encrypt_and_decrypt_data(self):
|
|
"""Test encrypt and decrypt utility functions."""
|
|
key = generate_encryption_key()
|
|
test_data = "sensitive data"
|
|
|
|
# Encrypt data
|
|
encrypted = encrypt_data(test_data, key)
|
|
assert isinstance(encrypted, bytes)
|
|
assert encrypted != test_data.encode()
|
|
|
|
# Decrypt data
|
|
decrypted = decrypt_data(encrypted, key)
|
|
assert decrypted == test_data
|
|
|
|
def test_encrypt_and_decrypt_dict_data(self):
|
|
"""Test encrypt and decrypt dictionary data."""
|
|
key = generate_encryption_key()
|
|
test_data = {"user": "john", "password": "secret123"}
|
|
|
|
# Encrypt data
|
|
encrypted = encrypt_data(test_data, key)
|
|
assert isinstance(encrypted, bytes)
|
|
|
|
# Decrypt data
|
|
decrypted = decrypt_data(encrypted, key)
|
|
assert decrypted == test_data
|
|
|
|
def test_encrypt_and_decrypt_binary_data(self):
|
|
"""Test encrypt and decrypt binary data."""
|
|
key = generate_encryption_key()
|
|
test_data = b"binary sensitive data"
|
|
|
|
# Encrypt data
|
|
encrypted = encrypt_data(test_data, key)
|
|
assert isinstance(encrypted, bytes)
|
|
assert encrypted != test_data
|
|
|
|
# Decrypt data
|
|
decrypted = decrypt_data(encrypted, key)
|
|
assert decrypted == test_data
|
|
|
|
def test_encrypt_with_invalid_key(self):
|
|
"""Test encrypt with invalid key."""
|
|
invalid_key = b"invalid_key"
|
|
test_data = "test data"
|
|
|
|
with pytest.raises(Exception):
|
|
encrypt_data(test_data, invalid_key)
|
|
|
|
def test_decrypt_with_invalid_key(self):
|
|
"""Test decrypt with invalid key."""
|
|
key = generate_encryption_key()
|
|
test_data = "test data"
|
|
encrypted = encrypt_data(test_data, key)
|
|
|
|
invalid_key = b"invalid_key"
|
|
with pytest.raises(Exception):
|
|
decrypt_data(encrypted, invalid_key)
|