Spaces:
Paused
Paused
File size: 3,629 Bytes
4a2ab42 291a9ec 4a2ab42 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | from cryptography.fernet import Fernet
from sqlalchemy.types import Text, TypeDecorator
from core.config import settings
from core.logging import logger
class VersionedEncryptedString(TypeDecorator):
"""
SQLAlchemy TypeDecorator that encrypts data with key versioning for migration support.
Supports multiple encryption keys and automatic decryption with fallback.
"""
impl = Text
cache_ok = True
# Key versions for migration support
KEY_VERSIONS = {
"v1": settings.FIELD_ENCRYPTION_KEY, # Current key
# Add older keys here for migration: "v0": "old_key_here"
}
def __init__(self, key=None, **kwargs):
super().__init__(**kwargs)
self._current_version = "v1"
self._fernet_instances = {}
# Initialize Fernet instances for all key versions
for version, key_value in self.KEY_VERSIONS.items():
key_bytes = key_value.encode() if isinstance(key_value, str) else key_value
self._fernet_instances[version] = Fernet(key_bytes)
# Use current version for encryption
self._current_fernet = self._fernet_instances[self._current_version]
def process_bind_param(self, value, dialect):
"""Encrypt value before saving with version prefix"""
if value is None:
return None
if not isinstance(value, str):
value = str(value)
try:
encrypted = self._current_fernet.encrypt(value.encode()).decode()
# Prefix with version for future migration support
return f"{self._current_version}:{encrypted}"
except Exception as e:
logger.error(f"Encryption failed: {e}")
raise
def process_result_value(self, value, dialect):
"""Decrypt value after retrieving with version-aware fallback"""
if value is None:
return None
# Handle versioned encrypted data
if ":" in value and len(value.split(":", 1)) == 2:
version, encrypted_data = value.split(":", 1)
if version in self._fernet_instances:
try:
return self._fernet_instances[version].decrypt(encrypted_data.encode()).decode()
except Exception as e:
logger.error(f"Decryption failed for {version}: {e}")
# Continue to fallback logic
else:
logger.warning(f"Unknown encryption version: {version}")
# Fallback: try all known keys (for legacy data without version prefix)
for version, fernet in self._fernet_instances.items():
try:
return fernet.decrypt(value.encode()).decode()
except Exception:
continue
# Final fallback: return raw value (might be unencrypted legacy data)
logger.error(f"All decryption attempts failed for value starting with: {value[:50]}...")
return value
# Backward compatibility alias
class EncryptedString(VersionedEncryptedString):
"""Backward compatibility alias for VersionedEncryptedString"""
cache_ok = True
# Helper methods
def encrypt_value(value: str) -> str:
if value is None:
return None
key = settings.FIELD_ENCRYPTION_KEY
if isinstance(key, str):
key = key.encode()
f = Fernet(key)
return f.encrypt(value.encode()).decode()
def decrypt_value(token: str) -> str:
if token is None:
return None
key = settings.FIELD_ENCRYPTION_KEY
if isinstance(key, str):
key = key.encode()
f = Fernet(key)
return f.decrypt(token.encode()).decode()
|