75 lines
1.9 KiB
Python
75 lines
1.9 KiB
Python
|
"""Encryption utilities for sensitive data like tokens."""
|
||
|
|
|
||
|
|
import logging
|
||
|
|
import os
|
||
|
|
|
||
|
|
from cryptography.fernet import Fernet, InvalidToken
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class EncryptionKeyMissingError(ValueError):
|
||
|
|
"""Raised when TOKEN_ENCRYPTION_KEY environment variable is not set."""
|
||
|
|
|
||
|
|
|
||
|
|
def _get_encryption_key() -> bytes:
|
||
|
|
"""Get or derive the encryption key from environment variable.
|
||
|
|
|
||
|
|
Uses TOKEN_ENCRYPTION_KEY env var if set (must be 32 url-safe base64 bytes),
|
||
|
|
otherwise derives a key from LANGSMITH_API_KEY using SHA256.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
32-byte Fernet-compatible key
|
||
|
|
|
||
|
|
Raises:
|
||
|
|
EncryptionKeyMissingError: If TOKEN_ENCRYPTION_KEY is not set
|
||
|
|
"""
|
||
|
|
explicit_key = os.environ.get("TOKEN_ENCRYPTION_KEY")
|
||
|
|
if not explicit_key:
|
||
|
|
raise EncryptionKeyMissingError
|
||
|
|
|
||
|
|
return explicit_key.encode()
|
||
|
|
|
||
|
|
|
||
|
|
def encrypt_token(token: str) -> str:
|
||
|
|
"""Encrypt a token for safe storage.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
token: The plaintext token to encrypt
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Base64-encoded encrypted token
|
||
|
|
"""
|
||
|
|
if not token:
|
||
|
|
return ""
|
||
|
|
|
||
|
|
key = _get_encryption_key()
|
||
|
|
f = Fernet(key)
|
||
|
|
encrypted = f.encrypt(token.encode())
|
||
|
|
return encrypted.decode()
|
||
|
|
|
||
|
|
|
||
|
|
def decrypt_token(encrypted_token: str) -> str:
|
||
|
|
"""Decrypt an encrypted token.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
encrypted_token: The base64-encoded encrypted token
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
The plaintext token, or empty string if decryption fails
|
||
|
|
"""
|
||
|
|
if not encrypted_token:
|
||
|
|
return ""
|
||
|
|
|
||
|
|
try:
|
||
|
|
key = _get_encryption_key()
|
||
|
|
f = Fernet(key)
|
||
|
|
decrypted = f.decrypt(encrypted_token.encode())
|
||
|
|
return decrypted.decode()
|
||
|
|
except InvalidToken:
|
||
|
|
logger.warning("Failed to decrypt token: invalid token")
|
||
|
|
return ""
|
||
|
|
except EncryptionKeyMissingError:
|
||
|
|
logger.warning("Failed to decrypt token: encryption key not set")
|
||
|
|
return ""
|