Files
xiaobaifupan/next/backend/security/cipher.py
T

51 lines
1.8 KiB
Python

from __future__ import annotations
import os
from dataclasses import dataclass
from cryptography.fernet import Fernet, InvalidToken
from backend.bootstrap.settings import ConfigurationError, Settings
@dataclass(frozen=True, slots=True)
class SecretCipher:
_fernet: Fernet
@classmethod
def from_key(cls, key: str) -> SecretCipher:
try:
return cls(Fernet(key.encode("ascii")))
except (ValueError, UnicodeEncodeError) as exc:
raise ConfigurationError("APP_ENCRYPTION_KEY is not a valid Fernet key") from exc
def encrypt(self, value: str) -> str:
return self._fernet.encrypt(value.encode("utf-8")).decode("ascii")
def decrypt(self, value: str) -> str:
try:
return self._fernet.decrypt(value.encode("ascii")).decode("utf-8")
except (InvalidToken, UnicodeDecodeError, UnicodeEncodeError) as exc:
raise ValueError("Encrypted value cannot be decrypted") from exc
def load_or_create_cipher(settings: Settings) -> SecretCipher:
if settings.encryption_key is not None:
return SecretCipher.from_key(settings.encryption_key)
if settings.environment == "production":
raise ConfigurationError("APP_ENCRYPTION_KEY is required in production")
key_file = settings.data_directory / ".encryption.key"
key_file.parent.mkdir(parents=True, exist_ok=True)
if key_file.exists():
return SecretCipher.from_key(key_file.read_text(encoding="ascii").strip())
key = Fernet.generate_key().decode("ascii")
try:
with key_file.open("x", encoding="ascii") as handle:
handle.write(key)
os.chmod(key_file, 0o600)
except FileExistsError:
key = key_file.read_text(encoding="ascii").strip()
return SecretCipher.from_key(key)