From 7ccaa654c2f44501e8cd0a3a0f388400dee6f797 Mon Sep 17 00:00:00 2001 From: Dillon Walls Date: Thu, 12 Sep 2024 20:47:45 +0000 Subject: [PATCH] fixup! fixup! fixup! add support for mongodb Client Side Field Level Encryption (CSFLE) --- ming/encryption.py | 35 +++++++++++++++++++++---------- ming/tests/__init__.py | 38 ++++++++++++++++++++++++++++++++++ ming/tests/test_datastore.py | 15 ++++---------- ming/tests/test_declarative.py | 7 ++++--- 4 files changed, 70 insertions(+), 25 deletions(-) diff --git a/ming/encryption.py b/ming/encryption.py index 19711e4..eb9497f 100644 --- a/ming/encryption.py +++ b/ming/encryption.py @@ -1,12 +1,14 @@ 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 @@ -14,9 +16,20 @@ 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,6 +81,7 @@ 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] @@ -73,10 +89,7 @@ def make_data_key(cls): 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: @@ -84,11 +97,11 @@ def encr(cls, s: str | None, _first_attempt=True, provider='local') -> bytes | N 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) diff --git a/ming/tests/__init__.py b/ming/tests/__init__.py index e69de29..f9f6770 100644 --- a/ming/tests/__init__.py +++ b/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"]', + } + diff --git a/ming/tests/test_datastore.py b/ming/tests/test_datastore.py index 265ef65..bb9c04a 100644 --- a/ming/tests/test_datastore.py +++ b/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 diff --git a/ming/tests/test_declarative.py b/ming/tests/test_declarative.py index 4b48a3b..13f19d4 100644 --- a/ming/tests/test_declarative.py +++ b/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',