From 1e118db767043f87097959846a5fa37d5f42cab9 Mon Sep 17 00:00:00 2001
From: Dillon Walls <dillon.walls@slashdotmedia.com>
Date: Thu, 12 Sep 2024 17:03:56 -0400
Subject: [PATCH] fixup! fixup! fixup! fixup! add support for mongodb Client
 Side Field Level Encryption (CSFLE)

---
 ming/encryption.py             | 19 +++++++++++++++++
 ming/tests/__init__.py         | 37 +++++++++-------------------------
 ming/tests/test_datastore.py   | 17 ++++++++--------
 ming/tests/test_declarative.py |  4 ++--
 4 files changed, 40 insertions(+), 37 deletions(-)

diff --git a/ming/encryption.py b/ming/encryption.py
index eb9497f..13ea461 100644
--- a/ming/encryption.py
+++ b/ming/encryption.py
@@ -1,7 +1,11 @@
 from __future__ import annotations
 
+import base64
 from copy import deepcopy
+from contextlib import contextmanager
 import json
+import os
+import random
 from typing import TYPE_CHECKING, TypeVar, Generic
 
 from cachetools import cached, RRCache
@@ -16,6 +20,21 @@
     import ming.datastore
 
 
+@contextmanager
+def push_seed(seed):
+    rstate = random.getstate()
+    random.seed(seed)
+    try:
+        yield
+    finally:
+        random.setstate(rstate)
+
+
+def make_encryption_key(seed=__name__):
+    with push_seed(seed):
+        return base64.b64encode(os.urandom(96)).decode('ascii')
+
+
 class MingEncryptionError(Exception):
     pass
 
diff --git a/ming/tests/__init__.py b/ming/tests/__init__.py
index f9f6770..2dab4ef 100644
--- a/ming/tests/__init__.py
+++ b/ming/tests/__init__.py
@@ -1,35 +1,18 @@
-import base64
-import os
-import random
+from ming.encryption import make_encryption_key
 
-from contextlib import contextmanager
 
+class EncryptionConfigHelper:
 
-@contextmanager
-def push_seed(seed):
-    rstate = random.getstate()
-    random.seed(seed)
-    try:
-        yield
-    finally:
-        random.setstate(rstate)
-
-
-class MingTestHelpers:
-
-    @staticmethod
-    def make_encryption_key(seed=__name__):
-        with push_seed(seed):
-            return base64.b64encode(os.urandom(96)).decode('ascii')
-
-    KEY_VAULT_NAMESPACE = 'encryption_test.coll_key_vault_test'
-    LOCAL_KEY = make_encryption_key('default local key')
+    LOCAL_KEY_VAULT_NAMESPACE = 'encryption_test.coll_key_vault_test'
+    LOCAL_KEY = make_encryption_key('test local key')
 
     @classmethod
-    def get_encrypt_cfg_local(cls,
-                              dbname: str,
-                              encryption_key = LOCAL_KEY,
-                              key_vault_namespace: str = KEY_VAULT_NAMESPACE):
+    def config_for_local_provider(
+        cls,
+        dbname: str,
+        encryption_key = LOCAL_KEY,
+        key_vault_namespace: str = LOCAL_KEY_VAULT_NAMESPACE
+    ):
         return {
             f'ming.{dbname}.encryption.kms_providers.local.key': encryption_key,
             f'ming.{dbname}.encryption.key_vault_namespace': key_vault_namespace,
diff --git a/ming/tests/test_datastore.py b/ming/tests/test_datastore.py
index bb9c04a..860ac87 100644
--- a/ming/tests/test_datastore.py
+++ b/ming/tests/test_datastore.py
@@ -9,9 +9,10 @@
 from ming import Session
 from ming import mim
 from ming import create_datastore, create_engine
-from ming.datastore import Engine
+from ming.encryption import make_encryption_key
 from ming.exc import MingConfigError
-from ming.tests import MingTestHelpers, make_encryption_key
+from ming.datastore import Engine
+from ming.tests import EncryptionConfigHelper
 
 
 class DummyConnection:
@@ -183,14 +184,14 @@ def test_configure_encryption(self):
             'ming.main.uri': 'mongodb://localhost:27017/test_db',
             'ming.main.replicaSet': 'foobar',
             'ming.main.foo.bar': 'foobar',
-            **MingTestHelpers.get_encrypt_cfg_local('main', encryption_key)
+            **EncryptionConfigHelper.config_for_local_provider('main', encryption_key)
         })
         session = Session.by_name('main')
-        assert session.bind.conn is not None, session.bind.conn
-        assert session.bind.db is not None, session.bind.db
-        assert session.bind.encryption.kms_providers == {'local': {'key': encryption_key}}, session.bind.encryption.kms_providers
-        assert session.bind.encryption.key_vault_namespace == 'encryption.collectionName', session.bind.encryption.key_vault_namespace
-        assert session.bind.encryption.provider_options == {'local': {'key_alt_names': ['datakeyName']}}, session.bind.encryption.provider_options
+        assert session.bind.conn is not None
+        assert session.bind.db is not None
+        assert session.bind.encryption.kms_providers == {'local': {'key': encryption_key}}
+        assert session.bind.encryption.key_vault_namespace == EncryptionConfigHelper.LOCAL_KEY_VAULT_NAMESPACE
+        assert session.bind.encryption.provider_options == {'local': {'key_alt_names': ['datakeyName']}}
 
     def test_no_kwargs_with_bind(self):
         self.assertRaises(
diff --git a/ming/tests/test_declarative.py b/ming/tests/test_declarative.py
index 13f19d4..25e0c02 100644
--- a/ming/tests/test_declarative.py
+++ b/ming/tests/test_declarative.py
@@ -9,13 +9,13 @@
 from ming.base import Cursor
 from ming.datastore import create_datastore
 from ming.declarative import Document
-from ming.encryption import EncryptionConfig, DecryptedField
+from ming.encryption import EncryptionConfig, DecryptedField, make_encryption_key
 from ming.metadata import Field, Index
 from ming import schema as S
 from ming.odm.odmsession import ODMSession, ThreadLocalODMSession
 from ming.session import Session
 from ming.exc import MingException
-from ming.tests import make_encryption_key
+
 
 def mock_datastore():
     def mock_collection():