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

…ion (CSFLE)

ming/datastore.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import time
 import logging
 from threading import Lock

ming/encryption.py

@@ -20,13 +20,16 @@ def __init__(self, encryption_config):
  self._encryption_config = encryption_config
 
  @property
- def kms_providers(self) -> str:
+ def kms_providers(self) -> dict:
  return self._encryption_config.get('kms_providers')
 
  @property
- def provider_options(self) -> str:
+ def provider_options(self) -> dict:
  return self._encryption_config.get('provider_options')
 
+ def key_alt_name(self, provider='local') -> str:
+ return self.provider_options.get(provider)['key_alt_names'][0]
+
  @property
  def key_vault_namespace(self) -> str:
  return self._encryption_config.get('key_vault_namespace')
@@ -56,33 +59,35 @@ class EncryptedDocumentMixin:
  @cached(RRCache(maxsize=99)) # needs to be per datastore, so we pass that as a param
  def encryptor(cls, ming_ds: ming.datastore.DataStore):
  conn: MongoClient = ming_ds.conn
- kms_providers = {"local": {"key": ming_ds.encryption_key}}
- encryption = ClientEncryption(kms_providers, ming_ds.encr_data_key_vault,
+ encryption = ClientEncryption(ming_ds.encryption.kms_providers, ming_ds.encryption.key_vault_namespace,
  conn, conn.codec_options)
  return encryption
 
  @classmethod
  def make_data_key(cls):
  ming_ds: ming.datastore.DataStore = cls.m.session.bind
  # index recommended by mongodb docs:
- key_vault_db_name, key_vault_coll_name = ming_ds.encr_data_key_vault.split('.')
+ 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}})
- cls.encryptor(ming_ds).create_data_key('local', key_alt_names=[ming_ds.encr_data_key_name])
+
+ 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'])
 
  @classmethod
- def encr(cls, s: str | None, _first_attempt=True) -> bytes | None:
+ 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=cls.ming_ds.encr_data_key_name)
+ key_alt_name=key_alt_name)
  except EncryptionError as e:
  if _first_attempt and 'not all keys requested were satisfied' in str(e):
  cls.make_data_key()

ming/tests/test_declarative.py

@@ -9,7 +9,7 @@
 from ming.base import Cursor
 from ming.datastore import create_datastore
 from ming.declarative import Document
-from ming.encryption import DecryptedField
+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
@@ -176,14 +176,29 @@ class TestDocumentEncryptionReal(TestCase):
  DATASTORE = f"mongodb://localhost/test_ming_TestDocumentReal_{os.getpid()}?serverSelectionTimeoutMS=100"
 
  def setUp(self):
- self.datastore = create_datastore(self.DATASTORE)
+ self.encryption_key = os.urandom(96)
+ # self.encryption_key = 'ODdlZGMzNjZlZWFmYTVlMDhhYWM0ZTBhNTQ5ZTE2YzQ3OWZmMzA5MDUxMDhhOTVlN2UyYTMzNzBkZDE5OGRhMg=='
+ encryption_config = EncryptionConfig({
+ 'kms_providers': {
+ 'local': {
+ 'key': self.encryption_key,
+ },
+ },
+ 'key_vault_namespace': 'encryption.__keyVault',
+ 'provider_options': {
+ 'local': {
+ 'key_alt_names': ['datakeyName'],
+ },
+ },
+ })
+ self.datastore = create_datastore(self.DATASTORE, encryption=encryption_config)
  self.session = Session(bind=self.datastore)
 
  class TestDoc(Document):
  class __mongometa__:
  name='test_doc'
  session = self.session
- indexes = [ ('a',) ]
+ indexes = [ ('name_encrypted',) ]
  _id = Field(S.Anything)
  name = DecryptedField(str, 'name_encrypted')
  name_encrypted = Field(S.Binary)
@@ -192,21 +207,19 @@ class __mongometa__:
 
  def tearDown(self):
  self.TestDoc.m.remove()
- # FIXME: teardown/ remove the encryption collection. likely in a different database
+ self.session.bind.conn.drop_database('encryption')
 
  def test_field(self):
- doc = self.TestDoc(dict(_id=1, a=1, b=dict(a=5)))
+ doc = self.TestDoc.make_encr(dict(_id=1, name='Jerome'))
  doc.m.save()
 
- self.assertEqual(doc.a, 1)
- self.assertEqual(doc.b, dict(a=5))
- doc.a = 5
- self.assertEqual(doc, dict(_id=1, a=5, b=dict(a=5)))
- del doc.a
- self.assertEqual(doc, dict(_id=1, b=dict(a=5)))
- self.assertRaises(AttributeError, getattr, doc, 'c')
- self.assertRaises(AttributeError, getattr, doc, 'a')
- self.assertEqual(self.session.count(self.TestDoc), 1)
+ self.assertEqual(doc.name, 'Jerome')
+ self.assertEqual(doc.name_encrypted, self.TestDoc.encr('Jerome'))
+ self.assertEqual(doc.name, self.TestDoc.decr(doc.name_encrypted))
+ doc.name = 'Jessie'
+ self.assertEqual(doc.name, 'Jessie')
+ self.assertEqual(doc.name_encrypted, self.TestDoc.encr('Jessie'))
+ self.assertEqual(doc.name, self.TestDoc.decr(doc.name_encrypted))
 
 
 class TestIndexes(TestCase):