INTPYTHON-527 Add Queryable Encryption support

django_mongodb_backend/base.py

@@ -1,9 +1,10 @@
 import contextlib
+import copy
 import os
 
 from django.core.exceptions import ImproperlyConfigured
 from django.db import DEFAULT_DB_ALIAS
-from django.db.backends.base.base import BaseDatabaseWrapper
+from django.db.backends.base.base import NO_DB_ALIAS, BaseDatabaseWrapper
 from django.db.backends.utils import debug_transaction
 from django.utils.asyncio import async_unsafe
 from django.utils.functional import cached_property
@@ -156,6 +157,9 @@ def _isnull_operator(a, b):
  def __init__(self, settings_dict, alias=DEFAULT_DB_ALIAS):
  super().__init__(settings_dict, alias=alias)
  self.session = None
+ # Cache the `settings_dict` in case we need to check for
+ # auto_encryption_opts later.
+ self.__dict__["_settings_dict"] = copy.deepcopy(settings_dict)
 
  def get_collection(self, name, **kwargs):
  collection = Collection(self.database, name, **kwargs)
@@ -287,3 +291,25 @@ def validate_no_broken_transaction(self):
  def get_database_version(self):
  """Return a tuple of the database's version."""
  return tuple(self.connection.server_info()["versionArray"])
+
+ @contextlib.contextmanager
+ def _nodb_cursor(self):
+ """
+ Returns a cursor from an unencrypted connection for operations
+ that do not support encryption.
+
+ Encryption is only supported on encrypted models.
+ """
+
+ # Remove auto_encryption_opts from OPTIONS
+ if self.settings_dict.get("OPTIONS", {}).get("auto_encryption_opts"):
+ self.settings_dict["OPTIONS"].pop("auto_encryption_opts")
+
+ # Create a new connection without OPTIONS["auto_encryption_opts": …]
+ conn = self.__class__({**self.settings_dict}, alias=NO_DB_ALIAS)
+
+ try:
+ with conn.cursor() as cursor:
+ yield cursor
+ finally:
+ conn.close()

django_mongodb_backend/encryption.py

@@ -0,0 +1,66 @@
+# Queryable Encryption helpers
+#
+# TODO: Decide if these helpers should even exist, and if so, find a permanent
+# place for them.
+
+from bson.binary import STANDARD
+from bson.codec_options import CodecOptions
+from pymongo.encryption import AutoEncryptionOpts, ClientEncryption
+
+
+def get_encrypted_client(auto_encryption_opts, encrypted_connection):
+ """
+ Returns a `ClientEncryption` instance for MongoDB Client-Side Field Level
+ Encryption (CSFLE) that can be used to create an encrypted collection.
+ """
+
+ key_vault_namespace = auto_encryption_opts._key_vault_namespace
+ kms_providers = auto_encryption_opts._kms_providers
+ codec_options = CodecOptions(uuid_representation=STANDARD)
+ return ClientEncryption(kms_providers, key_vault_namespace, encrypted_connection, codec_options)
+
+
+def get_auto_encryption_opts(crypt_shared_lib_path=None, kms_providers=None):
+ """
+ Returns an `AutoEncryptionOpts` instance for MongoDB Client-Side Field
+ Level Encryption (CSFLE) that can be used to create an encrypted connection.
+ """
+ key_vault_database_name = "encryption"
+ key_vault_collection_name = "__keyVault"
+ key_vault_namespace = f"{key_vault_database_name}.{key_vault_collection_name}"
+ return AutoEncryptionOpts(
+ key_vault_namespace=key_vault_namespace,
+ kms_providers=kms_providers,
+ crypt_shared_lib_path=crypt_shared_lib_path,
+ )
+
+
+def get_customer_master_key():
+ """
+ Returns a 96-byte local master key for use with MongoDB Client-Side Field Level
+ Encryption (CSFLE). For local testing purposes only. In production, use a secure KMS
+ like AWS, Azure, GCP, or KMIP.
+ Returns:
+ bytes: A 96-byte key.
+ """
+ # WARNING: This is a static key for testing only.
+ # Generate with: os.urandom(96)
+ return bytes.fromhex(
+ "000102030405060708090a0b0c0d0e0f"
+ "101112131415161718191a1b1c1d1e1f"
+ "202122232425262728292a2b2c2d2e2f"
+ "303132333435363738393a3b3c3d3e3f"
+ "404142434445464748494a4b4c4d4e4f"
+ "505152535455565758595a5b5c5d5e5f"
+ )
+
+
+def get_kms_providers():
+ """
+ Return supported KMS providers for MongoDB Client-Side Field Level Encryption (CSFLE).
+ """
+ return {
+ "local": {
+ "key": get_customer_master_key(),
+ },
+ }

django_mongodb_backend/features.py

@@ -624,3 +624,18 @@ def supports_transactions(self):
  hello = client.command("hello")
  # a replica set or a sharded cluster
  return "setName" in hello or hello.get("msg") == "isdbgrid"
+
+ @cached_property
+ def supports_encryption(self):
+ """
+ Encryption is supported if the server is Atlas or Enterprise
+ and is configured as a replica set or sharded cluster.
+ """
+ self.connection.ensure_connection()
+ client = self.connection.connection.admin
+ build_info = client.command("buildInfo")
+ is_enterprise = "enterprise" in build_info.get("modules")
+ # `supports_transactions` already checks if the server is a
+ # replica set or sharded cluster.
+ is_not_single = self.supports_transactions
+ return is_enterprise and is_not_single

django_mongodb_backend/fields/__init__.py

@@ -3,6 +3,7 @@
 from .duration import register_duration_field
 from .embedded_model import EmbeddedModelField
 from .embedded_model_array import EmbeddedModelArrayField
+from .encryption import EncryptedCharField
 from .json import register_json_field
 from .objectid import ObjectIdField
 
@@ -11,6 +12,7 @@
  "ArrayField",
  "EmbeddedModelArrayField",
  "EmbeddedModelField",
+ "EncryptedCharField",
  "ObjectIdAutoField",
  "ObjectIdField",
 ]

django_mongodb_backend/fields/encryption.py

@@ -0,0 +1,7 @@
+from django.db import models
+
+
+class EncryptedCharField(models.CharField):
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.encrypted = True

django_mongodb_backend/models.py

@@ -14,3 +14,26 @@ def delete(self, *args, **kwargs):
 
  def save(self, *args, **kwargs):
  raise NotSupportedError("EmbeddedModels cannot be saved.")
+
+
+class EncryptedModelBase(models.base.ModelBase):
+ def __new__(cls, name, bases, attrs, **kwargs):
+ new_class = super().__new__(cls, name, bases, attrs, **kwargs)
+
+ # Build a map of encrypted fields
+ encrypted_fields = {
+ "fields": {
+ field.name: field.__class__.__name__
+ for field in new_class._meta.fields
+ if getattr(field, "encrypted", False)
+ }
+ }
+
+ # Store it as a class-level attribute
+ new_class.encrypted_fields_map = encrypted_fields
+ return new_class
+
+
+class EncryptedModel(models.Model, metaclass=EncryptedModelBase):
+ class Meta:
+ abstract = True

django_mongodb_backend/schema.py

@@ -1,10 +1,13 @@
+import contextlib
+
 from django.db.backends.base.schema import BaseDatabaseSchemaEditor
 from django.db.models import Index, UniqueConstraint
+from pymongo.encryption import EncryptedCollectionError
 from pymongo.operations import SearchIndexModel
 
-from django_mongodb_backend.indexes import SearchIndex
-
+from .encryption import get_encrypted_client
 from .fields import EmbeddedModelField
+from .indexes import SearchIndex
 from .query import wrap_database_errors
 from .utils import OperationCollector
 
@@ -41,7 +44,7 @@ def get_database(self):
  @wrap_database_errors
  @ignore_embedded_models
  def create_model(self, model):
- self.get_database().create_collection(model._meta.db_table)
+ self._create_collection(model)
  self._create_model_indexes(model)
  # Make implicit M2M tables.
  for field in model._meta.local_many_to_many:
@@ -418,3 +421,45 @@ def _field_should_have_unique(self, field):
  db_type = field.db_type(self.connection)
  # The _id column is automatically unique.
  return db_type and field.unique and field.column != "_id"
+
+ def _supports_encryption(self, model):
+ """
+ Check for `supports_encryption` feature and `auto_encryption_opts`
+ and `embedded_fields_map`. If `supports_encryption` is True and
+ `auto_encryption_opts` is in the cached connection settings and
+ the model has an embedded_fields_map property, then encryption
+ is supported.
+ """
+ return (
+ self.connection.features.supports_encryption
+ and self.connection._settings_dict.get("OPTIONS", {}).get("auto_encryption_opts")
+ and hasattr(model, "encrypted_fields_map")
+ )
+
+ def _create_collection(self, model):
+ """
+ Create a collection or, if encryption is supported, create
+ an encrypted connection then use it to create an encrypted
+ client then use that to create an encrypted collection.
+ """
+
+ if self._supports_encryption(model):
+ auto_encryption_opts = self.connection._settings_dict.get("OPTIONS", {}).get(
+ "auto_encryption_opts"
+ )
+ # Use the cached settings dict to create a new connection
+ encrypted_connection = self.connection.get_new_connection(
+ self.connection._settings_dict
+ )
+ # Use the encrypted connection and auto_encryption_opts to create an encrypted client
+ encrypted_client = get_encrypted_client(auto_encryption_opts, encrypted_connection)
+
+ with contextlib.suppress(EncryptedCollectionError):
+ encrypted_client.create_encrypted_collection(
+ encrypted_connection[self.connection.database.name],
+ model._meta.db_table,
+ model.encrypted_fields_map,
+ "local", # TODO: KMS provider should be configurable
+ )
+ else:
+ self.get_database().create_collection(model._meta.db_table)

docs/source/topics/encrypted-models.rst

@@ -0,0 +1,22 @@
+Encrypted models
+================
+
+``EncryptedCharField``
+----------------------
+
+The basics
+~~~~~~~~~~
+
+Let's consider this example::
+
+ from django.db import models
+
+ from django_mongodb_backend.fields import EncryptedCharField
+ from django_mongodb_backend.models import EncryptedModel
+
+
+ class Person(EncryptedModel):
+ ssn = EncryptedCharField("ssn", max_length=11)
+
+ def __str__(self):
+ return self.ssn

docs/source/topics/index.rst

@@ -10,4 +10,5 @@ know:
 
  cache
  embedded-models
+ encrypted-models
  known-issues

tests/encryption_/models.py

@@ -0,0 +1,9 @@
+from django_mongodb_backend.fields import EncryptedCharField
+from django_mongodb_backend.models import EncryptedModel
+
+
+class Person(EncryptedModel):
+ ssn = EncryptedCharField("ssn", max_length=11)
+
+ def __str__(self):
+ return self.ssn

tests/encryption_/tests.py

@@ -0,0 +1,30 @@
+from django.test import TestCase
+
+from .models import Person
+
+
+class EncryptedModelTests(TestCase):
+ @classmethod
+ def setUpTestData(cls):
+ cls.objs = [Person.objects.create()]
+
+ def test_encrypted_fields_map_on_class(self):
+ expected = {
+ "fields": {
+ "ssn": "EncryptedCharField",
+ }
+ }
+ self.assertEqual(Person.encrypted_fields_map, expected)
+
+ def test_encrypted_fields_map_on_instance(self):
+ instance = Person(ssn="123-45-6789")
+ expected = {
+ "fields": {
+ "ssn": "EncryptedCharField",
+ }
+ }
+ self.assertEqual(instance.encrypted_fields_map, expected)
+
+ def test_non_encrypted_fields_not_included(self):
+ encrypted_field_names = Person.encrypted_fields_map.keys()
+ self.assertNotIn("ssn", encrypted_field_names)