58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
import os
|
|
import base64
|
|
from cryptography.fernet import Fernet, InvalidToken
|
|
from cryptography.hazmat.primitives import hashes
|
|
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
|
|
|
|
|
|
class KeyEncryption:
|
|
def __init__(self, encryption_key: str | None = None):
|
|
self._fernet = self._create_fernet(encryption_key)
|
|
|
|
def _create_fernet(self, key: str | None) -> Fernet:
|
|
if not key:
|
|
key = os.getenv("API_KEY_ENCRYPTION_KEY", "")
|
|
|
|
if not key:
|
|
return Fernet(Fernet.generate_key())
|
|
|
|
try:
|
|
return Fernet(key.encode())
|
|
except Exception:
|
|
kdf = PBKDF2HMAC(
|
|
algorithm=hashes.SHA256(),
|
|
length=32,
|
|
salt=b"geo-platform-salt",
|
|
iterations=100000,
|
|
)
|
|
derived_key = base64.urlsafe_b64encode(kdf.derive(key.encode()))
|
|
return Fernet(derived_key)
|
|
|
|
def encrypt(self, plaintext: str) -> str:
|
|
if not plaintext:
|
|
raise ValueError("Cannot encrypt empty string")
|
|
return self._fernet.encrypt(plaintext.encode()).decode()
|
|
|
|
def decrypt(self, ciphertext: str) -> str:
|
|
if not ciphertext:
|
|
raise ValueError("Cannot decrypt empty string")
|
|
try:
|
|
return self._fernet.decrypt(ciphertext.encode()).decode()
|
|
except InvalidToken:
|
|
raise ValueError("Invalid ciphertext or key")
|
|
|
|
|
|
_key_encryption: KeyEncryption | None = None
|
|
|
|
|
|
def get_key_encryption() -> KeyEncryption:
|
|
global _key_encryption
|
|
if _key_encryption is None:
|
|
_key_encryption = KeyEncryption()
|
|
return _key_encryption
|
|
|
|
|
|
def reset_key_encryption() -> None:
|
|
global _key_encryption
|
|
_key_encryption = None
|