INTPYTHON-527 Add Queryable Encryption support

Author: aclark4life

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_/__init__.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_/models.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)