fixup! fixup! fixup! fixup! fixup! add support for mongodb Client Sid…

…e Field Level Encryption (CSFLE)

ming/declarative.pyi

@@ -1,6 +1,9 @@
 from typing import TypeVar, Mapping, Any
 
+from pymongo.encryption import ClientEncryption
+
 from ming.base import Object
+from ming.datastore import DataStore
 from ming.metadata import Manager
 
 M = TypeVar('M')
@@ -16,6 +19,32 @@ class Document(Object):
  @classmethod
  def make(cls, data, allow_extra=False, strip_extra=True) -> Document: ...
 
+ # Encryption-Related fields:
+
+ @classmethod
+ def make_encr(cls, data: dict) -> Document: ...
+
+ @classmethod
+ def encr(cls, s: str | None, _first_attempt=True, provider='local') -> bytes | None: ...
+
+ @classmethod
+ def decr(cls, b: bytes | None) -> str | None: ...
+
+ @classmethod
+ def encryptor(cls, ming_ds: DataStore) -> ClientEncryption: ...
+
+ @classmethod
+ def make_data_key(cls): ...
+
+ @classmethod
+ def decrypted_field_names(cls) -> list[str]: ...
+
+ @classmethod
+ def encrypted_field_names(cls) -> list[str]: ...
+
+ @classmethod
+ def encrypt_some_fields(cls, data: dict) -> dict: ...
+
  # ...
  # class __mongometa__:
  # name: Any = ...

ming/encryption.py

@@ -1,11 +1,7 @@
 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
@@ -20,21 +16,6 @@
  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
 
@@ -89,7 +70,7 @@ class EncryptedDocumentMixin:
 
  @classmethod
  @cached(RRCache(maxsize=99)) # needs to be per datastore, so we pass that as a param
- def encryptor(cls, ming_ds: ming.datastore.DataStore):
+ def encryptor(cls, ming_ds: ming.datastore.DataStore) -> ClientEncryption:
  if not ming_ds.encryption:
  raise MingEncryptionError(f'No encryption settings found for {ming_ds}')
  conn: MongoClient = ming_ds.conn

ming/tests/__init__.py

@@ -1,8 +1,26 @@
-from ming.encryption import make_encryption_key
+import base64
+from contextlib import contextmanager
+import os
+import random
+
+
+@contextmanager
+def push_seed(seed):
+ rstate = random.getstate()
+ random.seed(seed)
+ try:
+ yield
+ finally:
+ random.setstate(rstate)
 
 
 class EncryptionConfigHelper:
 
+ @staticmethod
+ def make_encryption_key(seed=__name__):
+ with push_seed(seed):
+ return base64.b64encode(os.urandom(96)).decode('ascii')
+
  LOCAL_KEY_VAULT_NAMESPACE = 'encryption_test.coll_key_vault_test'
  LOCAL_KEY = make_encryption_key('test local key')
 

ming/tests/test_datastore.py

@@ -9,7 +9,6 @@
 from ming import Session
 from ming import mim
 from ming import create_datastore, create_engine
-from ming.encryption import make_encryption_key
 from ming.exc import MingConfigError
 from ming.datastore import Engine
 from ming.tests import EncryptionConfigHelper
@@ -179,7 +178,7 @@ def test_configure_optional_params(self):
  assert session.bind.db is not None
 
  def test_configure_encryption(self):
- encryption_key = make_encryption_key('foo')
+ encryption_key = EncryptionConfigHelper.make_encryption_key('foo')
  ming.configure(**{
  'ming.main.uri': 'mongodb://localhost:27017/test_db',
  'ming.main.replicaSet': 'foobar',

ming/tests/test_declarative.py

@@ -9,12 +9,13 @@
 from ming.base import Cursor
 from ming.datastore import create_datastore
 from ming.declarative import Document
-from ming.encryption import EncryptionConfig, DecryptedField, make_encryption_key
+from ming.encryption import EncryptionConfig, DecryptedField
 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 EncryptionConfigHelper
 
 
 def mock_datastore():
@@ -182,7 +183,7 @@ def setUp(self):
  encryption_config = EncryptionConfig({
  'kms_providers': {
  'local': {
- 'key': make_encryption_key(__name__),
+ 'key': EncryptionConfigHelper.make_encryption_key(__name__),
  },
  },
  'key_vault_namespace': 'encryption.__keyVault',