File size: 3,186 Bytes
f8b03c1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | """Module for encrypted."""
from __future__ import annotations
import base64
import logging
from typing import Any
from cryptography.fernet import Fernet, InvalidToken
from django.conf import settings
from django.db import models
logger = logging.getLogger(__name__)
class EncryptedTextField(models.TextField[str, str]):
prefix: str = "enc:v1:"
def __init__(self, *args: Any, **kwargs: Any) -> None:
super().__init__(*args, **kwargs)
def _resolve_keys(self) -> list[bytes]:
primary = getattr(settings, "CREDENTIAL_ENCRYPTION_KEY", None) or getattr(
settings, "SCRAPER_ENCRYPTION_KEY", None
)
previous = getattr(settings, "CREDENTIAL_ENCRYPTION_KEY_PREVIOUS", None) or getattr(
settings, "SCRAPER_ENCRYPTION_KEY_PREVIOUS", None
)
keys: list[bytes] = []
for k in (primary, previous):
if not k:
continue
if isinstance(k, str):
k = k.encode()
keys.append(k)
return keys
def _get_ciphers(self) -> list[Fernet]:
keys = self._resolve_keys()
if not keys:
raise RuntimeError("missing encryption key")
return [Fernet(k) for k in keys]
def _encrypt(self, plain_text: str) -> str:
cipher = self._get_ciphers()[0]
token = cipher.encrypt(plain_text.encode())
return f"{self.prefix}{base64.urlsafe_b64encode(token).decode()}"
def _decrypt(self, encrypted_value: str) -> str:
token_b64 = encrypted_value[len(self.prefix) :]
token = base64.urlsafe_b64decode(token_b64.encode())
last_error: Exception | None = None
for cipher in self._get_ciphers():
try:
decrypted: bytes = cipher.decrypt(token)
return decrypted.decode()
except Exception as e:
logger.debug("解密尝试失败，继续下一个密钥: %s", e)
last_error = e
continue
if last_error:
raise last_error
raise InvalidToken()
def get_prep_value(self, value: Any) -> Any:
plain_value = super().get_prep_value(value)
if plain_value is None or plain_value == "":
return plain_value
if isinstance(plain_value, str) and plain_value.startswith(self.prefix):
return plain_value
if not isinstance(plain_value, str):
plain_value = str(plain_value)
return self._encrypt(plain_value)
def to_python(self, value: Any) -> Any:
if value is None or value == "":
return value
if not isinstance(value, str):
value = str(value)
if not value.startswith(self.prefix):
return value
try:
return self._decrypt(value)
except (InvalidToken, ValueError):
if getattr(settings, "DEBUG", False):
logger.warning("解密失败，返回原始加密值（DEBUG 模式）: %s...", value[:20])
return value
raise
def from_db_value(self, value: Any, expression: Any, connection: Any) -> Any:
return self.to_python(value)
|