47 lines
1.5 KiB
Python
47 lines
1.5 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from cryptography.fernet import Fernet, InvalidToken
|
||
|
|
from django.conf import settings
|
||
|
|
|
||
|
|
|
||
|
|
class AiKeyEncryptionError(RuntimeError):
|
||
|
|
"""Raised when AI provider key encryption/decryption cannot proceed."""
|
||
|
|
|
||
|
|
|
||
|
|
PREFIX = "fernet:"
|
||
|
|
|
||
|
|
|
||
|
|
def encrypt_api_key(raw_api_key: str) -> str:
|
||
|
|
api_key = str(raw_api_key or "").strip()
|
||
|
|
if not api_key:
|
||
|
|
raise AiKeyEncryptionError("AI api_key cannot be blank")
|
||
|
|
token = _fernet().encrypt(api_key.encode("utf-8")).decode("ascii")
|
||
|
|
return f"{PREFIX}{token}"
|
||
|
|
|
||
|
|
|
||
|
|
def decrypt_api_key(encrypted_api_key: str) -> str:
|
||
|
|
encrypted = str(encrypted_api_key or "")
|
||
|
|
if not encrypted:
|
||
|
|
return ""
|
||
|
|
if not encrypted.startswith(PREFIX):
|
||
|
|
raise AiKeyEncryptionError("AI api_key is not stored in encrypted form")
|
||
|
|
token = encrypted[len(PREFIX) :].encode("ascii")
|
||
|
|
try:
|
||
|
|
return _fernet().decrypt(token).decode("utf-8")
|
||
|
|
except InvalidToken as exc:
|
||
|
|
raise AiKeyEncryptionError("AI api_key cannot be decrypted") from exc
|
||
|
|
|
||
|
|
|
||
|
|
def mask_secret(encrypted_api_key: str) -> str:
|
||
|
|
return "********" if encrypted_api_key else ""
|
||
|
|
|
||
|
|
|
||
|
|
def _fernet() -> Fernet:
|
||
|
|
key = getattr(settings, "AI_KEY_ENCRYPTION_KEY", "")
|
||
|
|
if not key:
|
||
|
|
raise AiKeyEncryptionError("AI_KEY_ENCRYPTION_KEY is not configured")
|
||
|
|
try:
|
||
|
|
return Fernet(str(key).encode("ascii"))
|
||
|
|
except (TypeError, ValueError) as exc:
|
||
|
|
raise AiKeyEncryptionError("AI_KEY_ENCRYPTION_KEY is invalid") from exc
|