Prevent showencryptedfieldsmap from creating data keys

Author: timgraham

django_mongodb_backend/management/commands/showencryptedfieldsmap.py

@@ -30,6 +30,6 @@ def handle(self, *args, **options):
             for app_config in apps.get_app_configs():
                 for model in router.get_migratable_models(app_config, db):
                     if model_has_encrypted_fields(model):
-                        fields = editor._get_encrypted_fields(model)
+                        fields = editor._get_encrypted_fields(model, create_data_keys=False)
                         encrypted_fields_map[model._meta.db_table] = fields
         self.stdout.write(json_util.dumps(encrypted_fields_map, indent=2))

django_mongodb_backend/schema.py

@@ -477,7 +477,9 @@ def _create_collection(self, model):
             # Unencrypted path
             db.create_collection(db_table)
 
-    def _get_encrypted_fields(self, model, key_alt_name_prefix=None, path_prefix=None):
+    def _get_encrypted_fields(
+        self, model, *, key_alt_name_prefix=None, path_prefix=None, create_data_keys=True
+    ):
         """
         Return the encrypted fields map for the given model. The "prefix"
         arguments are used when this method is called recursively on embedded
@@ -488,12 +490,12 @@ def _get_encrypted_fields(self, model, key_alt_name_prefix=None, path_prefix=Non
         key_alt_name_prefix = key_alt_name_prefix or model._meta.db_table
         path_prefix = path_prefix or ""
         auto_encryption_opts = client._options.auto_encryption_opts
-        key_vault_db, key_vault_collection = auto_encryption_opts._key_vault_namespace.split(".", 1)
-        key_vault_collection = client[key_vault_db][key_vault_collection]
+        _, key_vault_collection = auto_encryption_opts._key_vault_namespace.split(".", 1)
+        key_vault = self.get_collection(key_vault_collection)
         # Create partial unique index on keyAltNames.
         # TODO: find a better place for this. It only needs to run once for an
         # application's lifetime.
-        key_vault_collection.create_index(
+        key_vault.create_index(
             "keyAltNames", unique=True, partialFilterExpression={"keyAltNames": {"$exists": True}}
         )
         # Select the KMS provider.
@@ -516,22 +518,29 @@ def _get_encrypted_fields(self, model, key_alt_name_prefix=None, path_prefix=Non
                     field.embedded_model,
                     key_alt_name_prefix=key_alt_name,
                     path_prefix=path,
+                    create_data_keys=create_data_keys,
                 )
                 # An EmbeddedModelField may not have any encrypted fields.
                 if embedded_result:
                     field_list.extend(embedded_result["fields"])
                 continue
             # Populate data for encrypted field.
             if getattr(field, "encrypted", False):
-                data_key = key_vault_collection.find_one({"keyAltNames": key_alt_name})
-                if data_key:
-                    data_key = data_key["_id"]
-                else:
+                if create_data_keys:
                     data_key = connection.client_encryption.create_data_key(
                         kms_provider=kms_provider,
                         key_alt_names=[key_alt_name],
                         master_key=master_key,
                     )
+                else:
+                    data_key = key_vault.find_one({"keyAltNames": key_alt_name})
+                    if data_key:
+                        data_key = data_key["_id"]
+                    else:
+                        raise ImproperlyConfigured(
+                            f"Encryption key {key_alt_name} not found. Have "
+                            f"migrated the {model} model?"
+                        )
                 field_dict = {
                     "bsonType": field.db_type(connection),
                     "path": path,

django_mongodb_backend/utils.py

@@ -118,6 +118,7 @@ class OperationDebugWrapper:
         "create_indexes",
         "create_search_index",
         "drop",
+        "find_one",
         "index_information",
         "insert_many",
         "delete_many",

tests/encryption_/test_management.py

@@ -1,9 +1,12 @@
 from io import StringIO
 
 from bson import json_util
+from django.core.exceptions import ImproperlyConfigured
 from django.core.management import call_command
+from django.db import connections
 from django.test import modify_settings
 
+from .models import EncryptionKey
 from .test_base import EncryptionTestCase
 
 
@@ -96,7 +99,7 @@ class CommandTests(EncryptionTestCase):
 
     def _compare_output(self, expected, actual):
         for field in actual["fields"]:
-            field.pop("keyId", None)  # remove dynamic keyId
+            field.pop("keyId")  # remove dynamic keyId
         self.assertEqual(expected, actual)
 
     def test_show_encrypted_fields_map(self):
@@ -109,3 +112,20 @@ def test_show_encrypted_fields_map(self):
             with self.subTest(model=model_key):
                 self.assertIn(model_key, command_output)
                 self._compare_output(expected, command_output[model_key])
+
+    def test_missing_key(self):
+        test_key = "encryption__patient.patient_record.ssn"
+        msg = (
+            f"Encryption key {test_key} not found. Have migrated the "
+            "<class 'encryption_.models.PatientRecord'> model?"
+        )
+        EncryptionKey.objects.filter(key_alt_name=test_key).delete()
+        try:
+            with self.assertRaisesMessage(ImproperlyConfigured, msg):
+                call_command("showencryptedfieldsmap", "--database", "encrypted", verbosity=0)
+        finally:
+            # Replace the deleted key.
+            connections["encrypted"].client_encryption.create_data_key(
+                kms_provider="local",
+                key_alt_names=[test_key],
+            )

tests/encryption_/test_schema.py

@@ -96,16 +96,19 @@ def test_get_encrypted_fields_all_models(self):
         checks their encrypted fields map from the schema editor,
         and compares to expected BSON type & queries mapping.
         """
+        # Deleting all keys is only correct only if this test includes all
+        # test models. This test may not be needed since it's tested when the
+        # test runner migrates all models. If any subTest fails, the key vault
+        # will be left in an inconsistent state.
+        EncryptionKey.objects.all().delete()
         connection = connections["encrypted"]
-
         for model_name, expected in self.expected_map.items():
             with self.subTest(model=model_name):
                 model_class = getattr(models, model_name)
                 with connection.schema_editor() as editor:
-                    client = connection.connection
-                    encrypted_fields = editor._get_encrypted_fields(model_class, client)
+                    encrypted_fields = editor._get_encrypted_fields(model_class)
                     for field in encrypted_fields["fields"]:
-                        field.pop("keyId", None)  # Remove dynamic value
+                        field.pop("keyId")  # Remove dynamic value
                     self.assertEqual(encrypted_fields, expected)
 
     def test_key_creation_and_lookup(self):