149 lines
5.0 KiB
Python
149 lines
5.0 KiB
Python
import base64
|
|
import os
|
|
from typing import Optional
|
|
|
|
from cryptography.hazmat.backends import default_backend
|
|
from cryptography.hazmat.primitives import hashes
|
|
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
|
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
|
|
|
from letta.settings import settings
|
|
|
|
|
|
class CryptoUtils:
|
|
"""Utility class for AES-256-GCM encryption/decryption of sensitive data."""
|
|
|
|
# AES-256 requires 32 bytes key
|
|
KEY_SIZE = 32
|
|
# GCM standard IV size is 12 bytes (96 bits)
|
|
IV_SIZE = 12
|
|
# GCM tag size is 16 bytes (128 bits)
|
|
TAG_SIZE = 16
|
|
# Salt size for key derivation
|
|
SALT_SIZE = 16
|
|
|
|
@classmethod
|
|
def _derive_key(cls, master_key: str, salt: bytes) -> bytes:
|
|
"""Derive an AES key from the master key using PBKDF2."""
|
|
kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=cls.KEY_SIZE, salt=salt, iterations=100000, backend=default_backend())
|
|
return kdf.derive(master_key.encode())
|
|
|
|
@classmethod
|
|
def encrypt(cls, plaintext: str, master_key: Optional[str] = None) -> str:
|
|
"""
|
|
Encrypt a string using AES-256-GCM.
|
|
|
|
Args:
|
|
plaintext: The string to encrypt
|
|
master_key: Optional master key (defaults to settings.encryption_key)
|
|
|
|
Returns:
|
|
Base64 encoded string containing: salt + iv + ciphertext + tag
|
|
|
|
Raises:
|
|
ValueError: If no encryption key is configured
|
|
"""
|
|
if master_key is None:
|
|
master_key = settings.encryption_key
|
|
|
|
if not master_key:
|
|
raise ValueError(
|
|
"No encryption key configured. Please set the LETTA_ENCRYPTION_KEY environment variable (not fully supported yet for Letta v0.12.1 and below)."
|
|
)
|
|
|
|
# Generate random salt and IV
|
|
salt = os.urandom(cls.SALT_SIZE)
|
|
iv = os.urandom(cls.IV_SIZE)
|
|
|
|
# Derive key from master key
|
|
key = cls._derive_key(master_key, salt)
|
|
|
|
# Create cipher
|
|
cipher = Cipher(algorithms.AES(key), modes.GCM(iv), backend=default_backend())
|
|
encryptor = cipher.encryptor()
|
|
|
|
# Encrypt the plaintext
|
|
ciphertext = encryptor.update(plaintext.encode()) + encryptor.finalize()
|
|
|
|
# Get the authentication tag
|
|
tag = encryptor.tag
|
|
|
|
# Combine salt + iv + ciphertext + tag
|
|
encrypted_data = salt + iv + ciphertext + tag
|
|
|
|
# Return as base64 encoded string
|
|
return base64.b64encode(encrypted_data).decode("utf-8")
|
|
|
|
@classmethod
|
|
def decrypt(cls, encrypted: str, master_key: Optional[str] = None) -> str:
|
|
"""
|
|
Decrypt a string that was encrypted using AES-256-GCM.
|
|
|
|
Args:
|
|
encrypted: Base64 encoded encrypted string
|
|
master_key: Optional master key (defaults to settings.encryption_key)
|
|
|
|
Returns:
|
|
The decrypted plaintext string
|
|
|
|
Raises:
|
|
ValueError: If no encryption key is configured or decryption fails
|
|
"""
|
|
if master_key is None:
|
|
master_key = settings.encryption_key
|
|
|
|
if not master_key:
|
|
raise ValueError(
|
|
"No encryption key configured. Please set the LETTA_ENCRYPTION_KEY environment variable (not fully supported yet for Letta v0.12.1 and below)."
|
|
)
|
|
|
|
try:
|
|
# Decode from base64
|
|
encrypted_data = base64.b64decode(encrypted)
|
|
|
|
# Extract components
|
|
salt = encrypted_data[: cls.SALT_SIZE]
|
|
iv = encrypted_data[cls.SALT_SIZE : cls.SALT_SIZE + cls.IV_SIZE]
|
|
ciphertext = encrypted_data[cls.SALT_SIZE + cls.IV_SIZE : -cls.TAG_SIZE]
|
|
tag = encrypted_data[-cls.TAG_SIZE :]
|
|
|
|
# Derive key from master key
|
|
key = cls._derive_key(master_key, salt)
|
|
|
|
# Create cipher
|
|
cipher = Cipher(algorithms.AES(key), modes.GCM(iv, tag), backend=default_backend())
|
|
decryptor = cipher.decryptor()
|
|
|
|
# Decrypt the ciphertext
|
|
plaintext = decryptor.update(ciphertext) + decryptor.finalize()
|
|
|
|
return plaintext.decode("utf-8")
|
|
|
|
except Exception as e:
|
|
raise ValueError(f"Failed to decrypt data: {str(e)}")
|
|
|
|
@classmethod
|
|
def is_encrypted(cls, value: str) -> bool:
|
|
"""
|
|
Check if a string appears to be encrypted (base64 encoded with correct size).
|
|
|
|
This is a heuristic check and may have false positives.
|
|
"""
|
|
try:
|
|
decoded = base64.b64decode(value)
|
|
# Check if length is consistent with our encryption format
|
|
# Minimum size: salt(16) + iv(12) + tag(16) + at least 1 byte of ciphertext
|
|
return len(decoded) >= cls.SALT_SIZE + cls.IV_SIZE + cls.TAG_SIZE + 1
|
|
except Exception:
|
|
return False
|
|
|
|
@classmethod
|
|
def is_encryption_available(cls) -> bool:
|
|
"""
|
|
Check if encryption is available (encryption key is configured).
|
|
|
|
Returns:
|
|
True if encryption key is configured, False otherwise
|
|
"""
|
|
return bool(settings.encryption_key)
|