diff --git a/ming/declarative.pyi b/ming/declarative.pyi index 87d026a..14088c2 100644 --- a/ming/declarative.pyi +++ b/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 = ... diff --git a/ming/encryption.py b/ming/encryption.py index 13ea461..9b16675 100644 --- a/ming/encryption.py +++ b/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 diff --git a/ming/tests/__init__.py b/ming/tests/__init__.py index 2dab4ef..e71f4aa 100644 --- a/ming/tests/__init__.py +++ b/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') diff --git a/ming/tests/test_datastore.py b/ming/tests/test_datastore.py index 860ac87..a58e1e0 100644 --- a/ming/tests/test_datastore.py +++ b/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', diff --git a/ming/tests/test_declarative.py b/ming/tests/test_declarative.py index 25e0c02..2a0cfe9 100644 --- a/ming/tests/test_declarative.py +++ b/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',