diff --git a/ming/datastore.py b/ming/datastore.py index d690406..5036aa0 100644 --- a/ming/datastore.py +++ b/ming/datastore.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import time import logging from threading import Lock diff --git a/ming/encryption.py b/ming/encryption.py index 1673023..19711e4 100644 --- a/ming/encryption.py +++ b/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,8 +59,7 @@ 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 @@ -65,24 +67,27 @@ def encryptor(cls, ming_ds: ming.datastore.DataStore): 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() diff --git a/ming/tests/test_declarative.py b/ming/tests/test_declarative.py index cc50aa7..4b48a3b 100644 --- a/ming/tests/test_declarative.py +++ b/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):