encryption.py 1.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243
  1. import os
  2. import base64
  3. from cryptography.fernet import Fernet
  4. from cryptography.hazmat.primitives import hashes
  5. from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
  6. def get_encryption_key() -> bytes:
  7. """Get or generate encryption key from environment variable"""
  8. key_string = os.environ.get('ENCRYPTION_KEY', 'default-encryption-key-change-in-production')
  9. # Derive a proper Fernet key from the string
  10. kdf = PBKDF2HMAC(
  11. algorithm=hashes.SHA256(),
  12. length=32,
  13. salt=b'aws-scanner-salt', # In production, use a proper random salt
  14. iterations=100000,
  15. )
  16. key = base64.urlsafe_b64encode(kdf.derive(key_string.encode()))
  17. return key
  18. def encrypt_value(value: str) -> str:
  19. """Encrypt a string value"""
  20. if not value:
  21. return None
  22. key = get_encryption_key()
  23. f = Fernet(key)
  24. encrypted = f.encrypt(value.encode('utf-8'))
  25. return base64.urlsafe_b64encode(encrypted).decode('utf-8')
  26. def decrypt_value(encrypted_value: str) -> str:
  27. """Decrypt an encrypted string value"""
  28. if not encrypted_value:
  29. return None
  30. key = get_encryption_key()
  31. f = Fernet(key)
  32. encrypted_bytes = base64.urlsafe_b64decode(encrypted_value.encode('utf-8'))
  33. decrypted = f.decrypt(encrypted_bytes)
  34. return decrypted.decode('utf-8')