fixup! fixup! fixup! add support for mongodb Client Side Field Level …

…Encryption (CSFLE)

ming/encryption.py

@@ -1,22 +1,35 @@
 from __future__ import annotations
 
+from copy import deepcopy
+import json
 from typing import TYPE_CHECKING, TypeVar, Generic
 
 from cachetools import cached, RRCache
 
 from pymongo.encryption import ClientEncryption, Algorithm
 from pymongo.errors import EncryptionError
-
+from pymongocrypt.errors import MongoCryptError
 
 if TYPE_CHECKING:
  from pymongo import MongoClient
  from ming import Document
  import ming.datastore
 
 
+class MingEncryptionError(Exception):
+ pass
+
+
 class EncryptionConfig:
 
- def __init__(self, encryption_config):
+ def __init__(self, encryption_config: dict):
+ encryption_config = deepcopy(encryption_config)
+
+ # when loaded from an ini via ming.configure, list values must be manually parsed
+ for provider, values in list(encryption_config['provider_options'].items()):
+ if 'key_alt_names' in values and not isinstance(values['key_alt_names'], list):
+ encryption_config['provider_options'][provider]['key_alt_names'] = json.loads(values['key_alt_names'])
+
  self._encryption_config = encryption_config
 
  @property
@@ -58,6 +71,8 @@ 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):
+ if not ming_ds.encryption:
+ raise MingEncryptionError(f'No encryption settings found for {ming_ds}')
  conn: MongoClient = ming_ds.conn
  encryption = ClientEncryption(ming_ds.encryption.kms_providers, ming_ds.encryption.key_vault_namespace,
  conn, conn.codec_options)
@@ -66,29 +81,27 @@ def encryptor(cls, ming_ds: ming.datastore.DataStore):
  @classmethod
  def make_data_key(cls):
  ming_ds: ming.datastore.DataStore = cls.m.session.bind
+ encryptor = cls.encryptor(ming_ds)
  # index recommended by mongodb docs:
  key_vault_db_name, key_vault_coll_name = ming_ds.encryption.key_vault_namespace.split('.')
  key_vault_coll = ming_ds.conn[key_vault_db_name][key_vault_coll_name]
  key_vault_coll.create_index("keyAltNames", unique=True,
  partialFilterExpression={"keyAltNames": {"$exists": True}})
 
  for provider, options in ming_ds.encryption.provider_options.items():
- cls.encryptor(ming_ds).create_data_key(provider, **options)
-
- # cls.encryptor(ming_ds).create_data_key('local', **ming_ds['provider_options']['local'])
- # cls.encryptor(ming_ds).create_data_key('aws', **ming_ds['provider_options']['aws'])
+ encryptor.create_data_key(provider, **options)
 
  @classmethod
  def encr(cls, s: str | None, _first_attempt=True, provider='local') -> bytes | None:
  if s is None:
  return None
  try:
  ming_ds: ming.datastore.DataStore = cls.m.session.bind
- key_alt_name = ming_ds.encryption.key_alt_name()
- return cls.encryptor(ming_ds).encrypt(s,
-  Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic,
-  key_alt_name=key_alt_name)
- except EncryptionError as e:
+ encryptor = cls.encryptor(ming_ds)
+ return encryptor.encrypt(s,
+ Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic,
+ key_alt_name=ming_ds.encryption.key_alt_name())
+ except (EncryptionError, MongoCryptError) as e:
  if _first_attempt and 'not all keys requested were satisfied' in str(e):
  cls.make_data_key()
  return cls.encr(s, _first_attempt=False)

ming/tests/__init__.py

@@ -0,0 +1,38 @@
+import base64
+import os
+import random
+
+from contextlib import contextmanager
+
+
+@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')
+
+ @classmethod
+ def get_encrypt_cfg_local(cls,
+ dbname: str,
+ encryption_key = LOCAL_KEY,
+ key_vault_namespace: str = KEY_VAULT_NAMESPACE):
+ return {
+ f'ming.{dbname}.encryption.kms_providers.local.key': encryption_key,
+ f'ming.{dbname}.encryption.key_vault_namespace': key_vault_namespace,
+ f'ming.{dbname}.encryption.provider_options.local.key_alt_names': '["datakeyName"]',
+ }
+

ming/tests/test_datastore.py

@@ -1,3 +1,4 @@
+import os
 import sys
 from unittest import TestCase, main
 
@@ -10,6 +11,7 @@
 from ming import create_datastore, create_engine
 from ming.datastore import Engine
 from ming.exc import MingConfigError
+from ming.tests import MingTestHelpers, make_encryption_key
 
 
 class DummyConnection:
@@ -176,21 +178,12 @@ def test_configure_optional_params(self):
  assert session.bind.db is not None
 
  def test_configure_encryption(self):
- # Generate a base64 encoded random string:
- # ```python
- # from base64 import b64encode; import secrets;
- # b64encode(secrets.token_hex().encode('utf8')).decode('utf8')
- # ```
- encryption_key = 'ODdlZGMzNjZlZWFmYTVlMDhhYWM0ZTBhNTQ5ZTE2YzQ3OWZmMzA5MDUxMDhhOTVlN2UyYTMzNzBkZDE5OGRhMg=='
+ encryption_key = make_encryption_key('foo')
  ming.configure(**{
  'ming.main.uri': 'mongodb://localhost:27017/test_db',
  'ming.main.replicaSet': 'foobar',
  'ming.main.foo.bar': 'foobar',
- 'ming.main.encryption.kms_providers.local.key': encryption_key,
- 'ming.main.encryption.key_vault_namespace': 'encryption.collectionName',
- 'ming.main.encryption.provider_options.local.key_alt_names': ['datakeyName'],
- # 'ming.main.encryption.provider_options.aws.key_alt_names': ['datakeyName'],
- # 'ming.main.encryption.provider_options.aws.master_key': ['datakeyName'],
+ **MingTestHelpers.get_encrypt_cfg_local('main', encryption_key)
  })
  session = Session.by_name('main')
  assert session.bind.conn is not None, session.bind.conn

ming/tests/test_declarative.py

@@ -15,6 +15,7 @@
 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():
@@ -176,12 +177,12 @@ class TestDocumentEncryptionReal(TestCase):
  DATASTORE = f"mongodb://localhost/test_ming_TestDocumentReal_{os.getpid()}?serverSelectionTimeoutMS=100"
 
  def setUp(self):
- self.encryption_key = os.urandom(96)
- # self.encryption_key = 'ODdlZGMzNjZlZWFmYTVlMDhhYWM0ZTBhNTQ5ZTE2YzQ3OWZmMzA5MDUxMDhhOTVlN2UyYTMzNzBkZDE5OGRhMg=='
+ # We could use variable_decode on flat list of config values here,
+ #  but that would force ming to rely on formencode.
  encryption_config = EncryptionConfig({
  'kms_providers': {
  'local': {
- 'key': self.encryption_key,
+ 'key': make_encryption_key(__name__),
  },
  },
  'key_vault_namespace': 'encryption.__keyVault',