1from __future__ import annotations
2
3import base64
4import json
5from functools import cache
6from typing import TYPE_CHECKING, Any
7
8try:
9 from cryptography.fernet import Fernet, InvalidToken, MultiFernet
10 from cryptography.hazmat.primitives import hashes
11 from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
12except ImportError:
13 Fernet = None # ty: ignore[invalid-assignment]
14 InvalidToken = None # ty: ignore[invalid-assignment]
15 MultiFernet = None # ty: ignore[invalid-assignment]
16 hashes = None # ty: ignore[invalid-assignment]
17 PBKDF2HMAC = None
18
19from plain.postgres.lookups import Exact, IsNull
20from plain.runtime import settings
21from plain.utils.encoding import force_bytes
22
23from plain import preflight
24
25from .base import NOT_PROVIDED
26from .json import JSONField
27from .text import TextField
28
29if TYPE_CHECKING:
30 from collections.abc import Callable, Sequence
31
32 from plain.postgres.connection import DatabaseConnection
33 from plain.postgres.lookups import Lookup, Transform
34 from plain.preflight.results import PreflightResult
35
36__all__ = [
37 "EncryptedJSONField",
38 "EncryptedTextField",
39]
40
41# Fixed salt for key derivation — changing this would invalidate all encrypted data.
42# This is not secret; it ensures the derived encryption key is distinct from
43# keys derived for other purposes (e.g., signing) even from the same SECRET_KEY.
44_KDF_SALT = b"plain.postgres.fields.encrypted"
45
46# Prefix for encrypted values in the database.
47# Makes encrypted data self-describing and distinguishable from plaintext.
48_ENCRYPTED_PREFIX = "$fernet$"
49
50
51def _derive_fernet_key(secret: str) -> bytes:
52 """Derive a Fernet-compatible key from an arbitrary secret string."""
53 if PBKDF2HMAC is None:
54 raise ImportError(
55 "The 'cryptography' package is required to use encrypted fields. "
56 "Install it with: pip install cryptography"
57 )
58 kdf = PBKDF2HMAC(
59 algorithm=hashes.SHA256(),
60 length=32,
61 salt=_KDF_SALT,
62 iterations=480_000,
63 )
64 return base64.urlsafe_b64encode(kdf.derive(force_bytes(secret)))
65
66
67@cache
68def _get_fernet(secret_key: str, fallbacks: tuple[str, ...]) -> MultiFernet:
69 """Build a MultiFernet from the given secret key and fallbacks.
70
71 The first key is used for encryption.
72 All keys are used for decryption, enabling key rotation.
73 Results are cached by (secret_key, fallbacks) so changing SECRET_KEY
74 (e.g. in tests) produces a new MultiFernet automatically.
75 """
76 keys = [_derive_fernet_key(secret_key)]
77 for fallback in fallbacks:
78 keys.append(_derive_fernet_key(fallback))
79 return MultiFernet([Fernet(k) for k in keys])
80
81
82def _encrypt(value: str) -> str:
83 """Encrypt a string and return a self-describing database value."""
84 if value == "":
85 return value
86 f = _get_fernet(settings.SECRET_KEY, tuple(settings.SECRET_KEY_FALLBACKS))
87 token = f.encrypt(force_bytes(value))
88 return _ENCRYPTED_PREFIX + token.decode("ascii")
89
90
91def _decrypt(value: str) -> str:
92 """Decrypt a self-describing database value back to a string.
93
94 Gracefully handles unencrypted values — if the value doesn't have
95 the encryption prefix, it's returned as-is. This supports gradual
96 migration from plaintext to encrypted fields.
97 """
98 if not value.startswith(_ENCRYPTED_PREFIX):
99 return value
100 token = value[len(_ENCRYPTED_PREFIX) :]
101 f = _get_fernet(settings.SECRET_KEY, tuple(settings.SECRET_KEY_FALLBACKS))
102 try:
103 return f.decrypt(token.encode("ascii")).decode("utf-8")
104 except InvalidToken:
105 raise ValueError(
106 "Could not decrypt field value. The SECRET_KEY (and SECRET_KEY_FALLBACKS) "
107 "may have changed since this data was encrypted."
108 )
109
110
111class EncryptedFieldMixin:
112 """Shared behavior for all encrypted fields.
113
114 Owns the lookup surface (isnull and exact only — ciphertext is
115 non-deterministic) and the preflight that blocks indexes and unique
116 constraints.
117
118 Must be used with Field as a co-base class.
119 """
120
121 # Type hints for attributes provided by Field (the required co-base class)
122 name: str
123 model: Any
124
125 # The complete lookup surface, replacing the base field's registry.
126 # isnull is obviously needed. exact is required so that `filter(field=None)`
127 # works — the ORM resolves "exact" first and then rewrites None to isnull.
128 # Exact lookups on non-None values will silently return no results (since
129 # ciphertext is non-deterministic), but blocking exact entirely would break
130 # the None/isnull path. The base classes are named directly — inheriting
131 # the concrete field's registrations would leak specialized lookups like
132 # JSONField's JSONExact, which compares against the jsonb 'null' literal
133 # and defeats the None→isnull rewrite. get_lookup()/get_transform() and
134 # registry consumers (e.g. unsupported-lookup error suggestions) all
135 # resolve through this one dict. A classmethod so both class-level and
136 # instance-level callers work.
137 @classmethod
138 def get_lookups(cls) -> dict[str, type[Lookup | Transform]]:
139 return {"exact": Exact, "isnull": IsNull}
140
141 def get_transform(self, name: str) -> Callable[..., Transform] | None:
142 # JSONField's get_transform falls back to KeyTransformFactory for any
143 # name — key transforms would operate on ciphertext, so block them.
144 return None
145
146 def preflight(self, **kwargs: Any) -> list[PreflightResult]:
147 errors: list[PreflightResult] = super().preflight(**kwargs) # ty: ignore[unresolved-attribute]
148 errors.extend(self._check_encrypted_constraints())
149 return errors
150
151 def _check_encrypted_constraints(self) -> list[PreflightResult]:
152 errors: list[PreflightResult] = []
153 if not hasattr(self, "model"):
154 return errors
155
156 field_name = self.name
157
158 for constraint in self.model.model_options.constraints:
159 constraint_fields = getattr(constraint, "fields", ())
160 if field_name in constraint_fields:
161 errors.append(
162 preflight.PreflightResult(
163 fix=(
164 f"'{self.model.__name__}.{field_name}' is an encrypted field "
165 f"and cannot be used in constraint '{constraint.name}'. "
166 "Encrypted values are non-deterministic."
167 ),
168 obj=self,
169 id="fields.encrypted_in_constraint",
170 )
171 )
172
173 for index in self.model.model_options.indexes:
174 index_fields = getattr(index, "fields", ())
175 # Strip ordering prefix (e.g., "-field_name" for descending)
176 stripped_fields = [f.lstrip("-") for f in index_fields]
177 if field_name in stripped_fields:
178 errors.append(
179 preflight.PreflightResult(
180 fix=(
181 f"'{self.model.__name__}.{field_name}' is an encrypted field "
182 f"and cannot be used in index '{index.name}'. "
183 "Encrypted values are non-deterministic."
184 ),
185 obj=self,
186 id="fields.encrypted_in_index",
187 )
188 )
189
190 return errors
191
192
193class EncryptedTextField[T: (str, str | None) = str](EncryptedFieldMixin, TextField[T]):
194 """A TextField that encrypts its value before storing in the database.
195
196 Values are encrypted using Fernet (AES-128-CBC + HMAC-SHA256) with a key
197 derived from SECRET_KEY. The database column is always ``text`` regardless
198 of max_length, since ciphertext length is unpredictable.
199
200 max_length is enforced on the plaintext value (validation), not on the
201 ciphertext stored in the database. Only ``default=""`` (with
202 ``required=False``) is accepted — empty strings are stored as plaintext
203 ``''``, so the empty value is the one default expressible as a column
204 DEFAULT; anything else would need ciphertext, which is non-deterministic.
205 """
206
207 only_empty_default = True
208
209 def __init__(
210 self,
211 *,
212 max_length: int | None = None,
213 required: bool = True,
214 allow_null: bool = False,
215 default: Any = NOT_PROVIDED,
216 validators: Sequence[Callable[..., Any]] = (),
217 ):
218 # Deliberately narrower than TextField: no `choices` — exact lookups
219 # on ciphertext are non-deterministic, so choice-based filtering would
220 # silently match nothing.
221 super().__init__(
222 max_length=max_length,
223 required=required,
224 allow_null=allow_null,
225 default=default,
226 validators=validators,
227 )
228
229 def get_db_prep_value(
230 self, value: Any, connection: DatabaseConnection, prepared: bool = False
231 ) -> Any:
232 value = super().get_db_prep_value(value, connection, prepared)
233 if value is None:
234 return value
235 return _encrypt(value)
236
237 def from_db_value(
238 self, value: Any, expression: Any, connection: DatabaseConnection
239 ) -> str | None:
240 if value is None:
241 return value
242 return _decrypt(value)
243
244
245class EncryptedJSONField(EncryptedFieldMixin, JSONField):
246 """A JSONField that encrypts its serialized value before storing in the database.
247
248 The JSON value is serialized to a string, encrypted, and stored as text.
249 On read, it's decrypted and deserialized back to a Python object.
250 """
251
252 db_type_sql = "text"
253 accepts_default = False
254
255 def __init__(
256 self,
257 *,
258 encoder: type[json.JSONEncoder] | None = None,
259 decoder: type[json.JSONDecoder] | None = None,
260 required: bool = True,
261 allow_null: bool = False,
262 validators: Sequence[Callable[..., Any]] = (),
263 ):
264 # Deliberately narrower than JSONField: no `default` — there is no
265 # empty plaintext value (even {} serializes to text that would need
266 # ciphertext, which is non-deterministic), so no literal column
267 # DEFAULT can be expressed.
268 super().__init__(
269 encoder=encoder,
270 decoder=decoder,
271 required=required,
272 allow_null=allow_null,
273 validators=validators,
274 )
275
276 def adapt_json_db_value(self, value: Any) -> Any:
277 # jsonb adaptation would emit jsonb — this column stores ciphertext.
278 if value is None:
279 return value
280 return _encrypt(json.dumps(value, cls=self.encoder))
281
282 def from_db_value(
283 self, value: Any, expression: Any, connection: DatabaseConnection
284 ) -> Any:
285 if value is None:
286 return value
287 decrypted = _decrypt(value)
288 try:
289 return json.loads(decrypted, cls=self.decoder)
290 except json.JSONDecodeError:
291 raise ValueError(
292 "Encrypted field contains data that is not valid JSON. "
293 "The stored value may be corrupt."
294 )