233 lines
8.8 KiB
Python
233 lines
8.8 KiB
Python
import base64
|
|
import json
|
|
import os
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from letta.helpers.crypto_utils import CryptoUtils
|
|
|
|
|
|
class TestCryptoUtils:
|
|
"""Test suite for CryptoUtils encryption/decryption functionality."""
|
|
|
|
# Mock master keys for testing
|
|
MOCK_KEY_1 = "test-master-key-1234567890abcdef"
|
|
MOCK_KEY_2 = "another-test-key-fedcba0987654321"
|
|
|
|
def test_encrypt_decrypt_roundtrip(self):
|
|
"""Test that encryption followed by decryption returns the original value."""
|
|
test_cases = [
|
|
"simple text",
|
|
"text with special chars: !@#$%^&*()",
|
|
"unicode text: 你好世界 🌍",
|
|
"very long text " * 1000,
|
|
'{"json": "data", "nested": {"key": "value"}}',
|
|
"", # Empty string
|
|
]
|
|
|
|
for plaintext in test_cases:
|
|
encrypted = CryptoUtils.encrypt(plaintext, self.MOCK_KEY_1)
|
|
assert encrypted != plaintext, f"Encryption failed for: {plaintext[:50]}"
|
|
# Encrypted value is base64 encoded
|
|
assert len(encrypted) > 0, "Encrypted value should not be empty"
|
|
|
|
decrypted = CryptoUtils.decrypt(encrypted, self.MOCK_KEY_1)
|
|
assert decrypted == plaintext, f"Roundtrip failed for: {plaintext[:50]}"
|
|
|
|
def test_encrypt_with_different_keys(self):
|
|
"""Test that different keys produce different ciphertexts."""
|
|
plaintext = "sensitive data"
|
|
|
|
encrypted1 = CryptoUtils.encrypt(plaintext, self.MOCK_KEY_1)
|
|
encrypted2 = CryptoUtils.encrypt(plaintext, self.MOCK_KEY_2)
|
|
|
|
# Different keys should produce different ciphertexts
|
|
assert encrypted1 != encrypted2
|
|
|
|
# Each should decrypt correctly with its own key
|
|
assert CryptoUtils.decrypt(encrypted1, self.MOCK_KEY_1) == plaintext
|
|
assert CryptoUtils.decrypt(encrypted2, self.MOCK_KEY_2) == plaintext
|
|
|
|
def test_decrypt_with_wrong_key_fails(self):
|
|
"""Test that decryption with wrong key raises an error."""
|
|
plaintext = "secret message"
|
|
encrypted = CryptoUtils.encrypt(plaintext, self.MOCK_KEY_1)
|
|
|
|
with pytest.raises(Exception): # Could be ValueError or cryptography exception
|
|
CryptoUtils.decrypt(encrypted, self.MOCK_KEY_2)
|
|
|
|
def test_encrypt_none_value(self):
|
|
"""Test handling of None values."""
|
|
# Encrypt None should raise TypeError (None has no encode method)
|
|
with pytest.raises((TypeError, AttributeError)):
|
|
CryptoUtils.encrypt(None, self.MOCK_KEY_1)
|
|
|
|
def test_decrypt_none_value(self):
|
|
"""Test that decrypting None raises an error."""
|
|
with pytest.raises(ValueError):
|
|
CryptoUtils.decrypt(None, self.MOCK_KEY_1)
|
|
|
|
def test_decrypt_empty_string(self):
|
|
"""Test that decrypting empty string raises an error."""
|
|
with pytest.raises(Exception): # base64 decode error
|
|
CryptoUtils.decrypt("", self.MOCK_KEY_1)
|
|
|
|
def test_decrypt_plaintext_value(self):
|
|
"""Test that decrypting non-encrypted value raises an error."""
|
|
plaintext = "not encrypted"
|
|
with pytest.raises(Exception): # Will fail base64 decode or decryption
|
|
CryptoUtils.decrypt(plaintext, self.MOCK_KEY_1)
|
|
|
|
def test_encrypted_format_structure(self):
|
|
"""Test the structure of encrypted values."""
|
|
plaintext = "test data"
|
|
encrypted = CryptoUtils.encrypt(plaintext, self.MOCK_KEY_1)
|
|
|
|
# Should be base64 encoded
|
|
encrypted_data = encrypted
|
|
|
|
# Should be valid base64
|
|
try:
|
|
decoded = base64.b64decode(encrypted_data)
|
|
assert len(decoded) > 0
|
|
except Exception as e:
|
|
pytest.fail(f"Invalid base64 encoding: {e}")
|
|
|
|
# Decoded data should contain salt, IV, tag, and ciphertext
|
|
# Total should be at least SALT_SIZE + IV_SIZE + TAG_SIZE bytes
|
|
min_size = CryptoUtils.SALT_SIZE + CryptoUtils.IV_SIZE + CryptoUtils.TAG_SIZE
|
|
assert len(decoded) >= min_size
|
|
|
|
def test_deterministic_with_same_salt(self):
|
|
"""Test that encryption is deterministic when using the same salt (for testing)."""
|
|
plaintext = "deterministic test"
|
|
|
|
# Note: In production, each encryption generates a random salt
|
|
# This test verifies the encryption mechanism itself
|
|
encrypted1 = CryptoUtils.encrypt(plaintext, self.MOCK_KEY_1)
|
|
encrypted2 = CryptoUtils.encrypt(plaintext, self.MOCK_KEY_1)
|
|
|
|
# Due to random salt, these should be different
|
|
assert encrypted1 != encrypted2
|
|
|
|
# But both should decrypt to the same value
|
|
assert CryptoUtils.decrypt(encrypted1, self.MOCK_KEY_1) == plaintext
|
|
assert CryptoUtils.decrypt(encrypted2, self.MOCK_KEY_1) == plaintext
|
|
|
|
def test_encrypt_uses_env_key_when_none_provided(self):
|
|
"""Test that encryption uses environment key when no key is provided."""
|
|
from letta.settings import settings
|
|
|
|
# Mock the settings to have an encryption key
|
|
original_key = settings.encryption_key
|
|
settings.encryption_key = "env-test-key-123"
|
|
|
|
try:
|
|
plaintext = "test with env key"
|
|
|
|
# Should use key from settings
|
|
encrypted = CryptoUtils.encrypt(plaintext)
|
|
assert len(encrypted) > 0
|
|
|
|
# Should decrypt with same key
|
|
decrypted = CryptoUtils.decrypt(encrypted)
|
|
assert decrypted == plaintext
|
|
finally:
|
|
# Restore original key
|
|
settings.encryption_key = original_key
|
|
|
|
def test_encrypt_without_key_raises_error(self):
|
|
"""Test that encryption without any key raises an error."""
|
|
from letta.settings import settings
|
|
|
|
# Mock settings to have no encryption key
|
|
original_key = settings.encryption_key
|
|
settings.encryption_key = None
|
|
|
|
try:
|
|
with pytest.raises(ValueError, match="No encryption key configured"):
|
|
CryptoUtils.encrypt("test data")
|
|
finally:
|
|
# Restore original key
|
|
settings.encryption_key = original_key
|
|
|
|
def test_large_data_encryption(self):
|
|
"""Test encryption of large data."""
|
|
# Create 10MB of data
|
|
large_data = "x" * (10 * 1024 * 1024)
|
|
|
|
encrypted = CryptoUtils.encrypt(large_data, self.MOCK_KEY_1)
|
|
assert len(encrypted) > 0
|
|
assert encrypted != large_data
|
|
|
|
decrypted = CryptoUtils.decrypt(encrypted, self.MOCK_KEY_1)
|
|
assert decrypted == large_data
|
|
|
|
def test_json_data_encryption(self):
|
|
"""Test encryption of JSON data."""
|
|
json_data = {
|
|
"user": "test_user",
|
|
"token": "secret_token_123",
|
|
"nested": {"api_key": "sk-1234567890", "headers": {"Authorization": "Bearer token"}},
|
|
}
|
|
|
|
json_str = json.dumps(json_data)
|
|
encrypted = CryptoUtils.encrypt(json_str, self.MOCK_KEY_1)
|
|
|
|
decrypted_str = CryptoUtils.decrypt(encrypted, self.MOCK_KEY_1)
|
|
decrypted_data = json.loads(decrypted_str)
|
|
|
|
assert decrypted_data == json_data
|
|
|
|
def test_invalid_encrypted_format(self):
|
|
"""Test handling of invalid encrypted data format."""
|
|
invalid_cases = [
|
|
"invalid-base64!@#", # Invalid base64
|
|
"dGVzdA==", # Valid base64 but too short for encrypted data
|
|
]
|
|
|
|
for invalid in invalid_cases:
|
|
with pytest.raises(Exception): # Could be various exceptions
|
|
CryptoUtils.decrypt(invalid, self.MOCK_KEY_1)
|
|
|
|
def test_key_derivation_consistency(self):
|
|
"""Test that key derivation is consistent."""
|
|
plaintext = "test key derivation"
|
|
|
|
# Multiple encryptions with same key should work
|
|
encrypted_values = []
|
|
for _ in range(5):
|
|
encrypted = CryptoUtils.encrypt(plaintext, self.MOCK_KEY_1)
|
|
encrypted_values.append(encrypted)
|
|
|
|
# All should decrypt correctly
|
|
for encrypted in encrypted_values:
|
|
assert CryptoUtils.decrypt(encrypted, self.MOCK_KEY_1) == plaintext
|
|
|
|
def test_special_characters_in_key(self):
|
|
"""Test encryption with keys containing special characters."""
|
|
special_key = "key-with-special-chars!@#$%^&*()_+"
|
|
plaintext = "test data"
|
|
|
|
encrypted = CryptoUtils.encrypt(plaintext, special_key)
|
|
decrypted = CryptoUtils.decrypt(encrypted, special_key)
|
|
|
|
assert decrypted == plaintext
|
|
|
|
def test_whitespace_handling(self):
|
|
"""Test encryption of strings with various whitespace."""
|
|
test_cases = [
|
|
" leading spaces",
|
|
"trailing spaces ",
|
|
" both sides ",
|
|
"multiple\n\nlines",
|
|
"\ttabs\there\t",
|
|
"mixed \t\n whitespace \r\n",
|
|
]
|
|
|
|
for plaintext in test_cases:
|
|
encrypted = CryptoUtils.encrypt(plaintext, self.MOCK_KEY_1)
|
|
decrypted = CryptoUtils.decrypt(encrypted, self.MOCK_KEY_1)
|
|
assert decrypted == plaintext, f"Whitespace handling failed for: {repr(plaintext)}"
|