Add queryable encryption config

Author: aclark4life

django_mongodb_backend/__init__.py

@@ -2,7 +2,7 @@
 
 # Check Django compatibility before other imports which may fail if the
 # wrong version of Django is installed.
-from .utils import check_django_compatability, parse_uri
+from .utils import check_django_compatability, get_auto_encryption_options, parse_uri
 
 check_django_compatability()
 
@@ -15,7 +15,7 @@
 from .lookups import register_lookups  # noqa: E402
 from .query import register_nodes  # noqa: E402
 
-__all__ = ["parse_uri"]
+__all__ = ["get_auto_encryption_options", "parse_uri"]
 
 register_aggregates()
 register_checks()

django_mongodb_backend/features.py

@@ -629,3 +629,13 @@ def supports_transactions(self):
             engine = client.command("serverStatus").get("storageEngine", {})
             return engine.get("name") == "wiredTiger"
         return False
+
+    @cached_property
+    def supports_queryable_encryption(self):
+        """
+        Queryable Encryption is available if the server is Atlas or Enterprise.
+        """
+        self.connection.ensure_connection()
+        client = self.connection.connection.admin
+        build_info = client.command("buildInfo")
+        return "enterprise" in build_info.get("modules")

django_mongodb_backend/utils.py

@@ -1,5 +1,8 @@
 import copy
+import os
 import time
+from pathlib import Path
+from urllib.parse import urlencode
 
 import django
 from django.conf import settings
@@ -8,6 +11,7 @@
 from django.utils.functional import SimpleLazyObject
 from django.utils.text import format_lazy
 from django.utils.version import get_version_tuple
+from pymongo.encryption_options import AutoEncryptionOpts
 from pymongo.uri_parser import parse_uri as pymongo_parse_uri
 
 
@@ -28,6 +32,62 @@ def check_django_compatability():
         )
 
 
+# Queryable Encryption-related functions based on helpers from Python Queryable
+# Encryption Tutorial
+# https://github.com/mongodb/docs/tree/master/source/includes/qe-tutorials/python/
+def _get_kms_provider_credentials(kms_provider_name):
+    """
+    "A KMS is a remote service that securely stores and manages your encryption keys."
+
+    Via https://www.mongodb.com/docs/manual/core/queryable-encryption/quick-start/
+
+    Here we check the provider name and return the appropriate credentials.
+    """
+    # TODO: Add support for other KMS providers.
+    if kms_provider_name == "local":
+        if not Path("./customer-master-key.txt").exists:
+            try:
+                path = "customer-master-key.txt"
+                file_bytes = os.urandom(96)
+                with Path.open(path, "wb") as f:
+                    f.write(file_bytes)
+            except Exception as e:
+                raise Exception(
+                    "Unable to write Customer Master Key to file due to the following error: "
+                ) from e
+
+        try:
+            path = "./customer-master-key.txt"
+            with Path.open(path, "rb") as f:
+                local_master_key = f.read()
+                if len(local_master_key) != 96:
+                    raise Exception("Expected the customer master key file to be 96 bytes.")
+                return {
+                    "local": {"key": local_master_key},
+                }
+        except Exception as e:
+            raise Exception(
+                "Unable to read Customer Master Key from file due to the following error: "
+            ) from e
+    else:
+        raise ValueError(
+            "Unrecognized value for kms_provider_name encountered while retrieving KMS credentials."
+        )
+
+
+def get_auto_encryption_options(kms_provider_name):
+    key_vault_database_name = "encryption"
+    key_vault_collection_name = "__keyVault"
+    key_vault_namespace = f"{key_vault_database_name}.{key_vault_collection_name}"
+    kms_provider_credentials = _get_kms_provider_credentials(kms_provider_name)
+    auto_encryption_opts = AutoEncryptionOpts(
+        kms_provider_credentials,
+        key_vault_namespace,
+        crypt_shared_lib_path=os.environ.get("SHARED_LIB_PATH"),
+    )
+    return urlencode(auto_encryption_opts)
+
+
 def parse_uri(uri, *, db_name=None, test=None):
     """
     Convert the given uri into a dictionary suitable for Django's DATABASES

tests/backend_/utils/test_parse_uri.py

@@ -1,10 +1,11 @@
 from unittest.mock import patch
+from urllib.parse import parse_qs
 
 import pymongo
 from django.core.exceptions import ImproperlyConfigured
 from django.test import SimpleTestCase
 
-from django_mongodb_backend import parse_uri
+from django_mongodb_backend import get_auto_encryption_options, parse_uri
 
 
 class ParseURITests(SimpleTestCase):
@@ -94,3 +95,10 @@ def test_invalid_credentials(self):
     def test_no_scheme(self):
         with self.assertRaisesMessage(pymongo.errors.InvalidURI, "Invalid URI scheme"):
             parse_uri("cluster0.example.mongodb.net")
+
+    def test_queryable_encryption_config(self):
+        auto_encryption_options = get_auto_encryption_options("local")
+        settings_dict = parse_uri(
+            f"mongodb://cluster0.example.mongodb.net/myDatabase{auto_encryption_options}"
+        )
+        self.assertEqual(settings_dict["OPTIONS"], parse_qs(auto_encryption_options))