MOTOR-689: Add async wrapper for pymongo.encryption.ClientEncryption

.travis.yml

@@ -10,6 +10,7 @@ services: mongodb
 
 install:
   - pip install tornado
+  - pip install pymongo[encryption]
 
 script: "python setup.py test"
 

doc/api-asyncio/asyncio_motor_client_encryption.rst

@@ -0,0 +1,7 @@
+:class:`~motor.motor_asyncio.AsyncIOMotorClientEncryption`
+==========================================================
+
+.. currentmodule:: motor.motor_asyncio
+
+.. autoclass:: AsyncIOMotorClientEncryption
+  :members:

doc/api-asyncio/index.rst

@@ -8,6 +8,7 @@ Motor asyncio API
     asyncio_motor_database
     asyncio_motor_collection
     asyncio_motor_change_stream
+    asyncio_motor_client_encryption
     cursors
     asyncio_gridfs
     aiohttp

doc/api-tornado/index.rst

@@ -8,6 +8,7 @@ Motor Tornado API
     motor_database
     motor_collection
     motor_change_stream
+    motor_client_encryption
     cursors
     gridfs
     web

doc/api-tornado/motor_client_encryption.rst

@@ -0,0 +1,7 @@
+:class:`~motor.motor_tornado.MotorClientEncryption`
+===================================================
+
+.. currentmodule:: motor.motor_tornado
+
+.. autoclass:: MotorClientEncryption
+  :members:

motor/core.py

@@ -34,6 +34,7 @@
 from pymongo.cursor import Cursor, RawBatchCursor, _QUERY_OPTIONS
 from pymongo.database import Database
 from pymongo.driver_info import DriverInfo
+from pymongo.encryption import ClientEncryption
 
 from . import version as motor_version
 from .metaprogramming import (AsyncCommand,
@@ -140,7 +141,7 @@ def __init__(self, *args, **kwargs):
 
         :Parameters:
           - `io_loop` (optional): Special event loop
-            instance to use instead of default
+            instance to use instead of default.
         """
         if 'io_loop' in kwargs:
             io_loop = kwargs.pop('io_loop')
@@ -1795,3 +1796,49 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         pass
+
+class AgnosticClientEncryption(AgnosticBase):
+    """Explicit client-side field level encryption."""
+
+    __motor_class_name__ = 'MotorClientEncryption'
+    __delegate_class__ = ClientEncryption
+
+    create_data_key = AsyncCommand(doc=create_data_key_doc)
+    encrypt = AsyncCommand()
+    decrypt = AsyncCommand()
+    close = AsyncCommand(doc=close_doc)
+
+    def __init__(self, kms_providers, key_vault_namespace, key_vault_client, codec_options, io_loop=None):
+        """Explicit client-side field level encryption.
+
+        Takes the same constructor arguments as
+        :class:`pymongo.encryption.ClientEncryption`, as well as:
+
+        :Parameters:
+          - `io_loop` (optional): Special event loop
+            instance to use instead of default.
+        """
+        if io_loop:
+            self._framework.check_event_loop(io_loop)
+        else:
+            io_loop = self._framework.get_event_loop()
+        sync_client = key_vault_client.delegate
+        delegate = self.__delegate_class__(kms_providers, key_vault_namespace, sync_client, codec_options)
+        super().__init__(delegate)
+        self.io_loop = io_loop
+
+    def get_io_loop(self):
+        return self.io_loop
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if self.delegate:
+            await self.close()
+
+    def __enter__(self):
+        raise RuntimeError('Use this encryption module in "async with", not "with"')
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass

motor/docstrings.py

@@ -1268,3 +1268,32 @@ async def coro():
 .. _$expr: https://docs.mongodb.com/manual/reference/operator/query/expr/
 .. _$where: https://docs.mongodb.com/manual/reference/operator/query/where/
 """
+
+create_data_key_doc = """Create and insert a new data key into the key vault collection.
+
+Takes the same constructors as 
+:class:`pymongo.encryption.ClientEncryption.create_data_key`,
+with only the following slight difference using async syntax:
+
+:Parameters:
+  - `key_alt_names` (optional): An optional list of string alternate
+    names used to reference a key. If a key is created with alternate
+    names, then encryption may refer to the key by the unique alternate
+    name instead of by ``key_id``. The following example shows creating
+    and referring to a data key by alternate name::
+
+      await client_encryption.create_data_key("local", keyAltNames=["name1"])
+      # reference the key with the alternate name
+      await client_encryption.encrypt("457-55-5462", keyAltName="name1",
+                                algorithm=Algorithm.Random)
+"""
+
+close_doc = """Release resources.
+
+Note that using this class in a with-statement will automatically call
+:meth:`close`::
+
+    async with AsyncIOMotorClientEncryption(...) as client_encryption:
+        encrypted = await client_encryption.encrypt(value, ...)
+        decrypted = await client_encryption.decrypt(encrypted)
+"""

motor/motor_asyncio.py

@@ -18,7 +18,7 @@
 from .frameworks import asyncio as asyncio_framework
 from .metaprogramming import create_class_with_framework
 
-__all__ = ['AsyncIOMotorClient']
+__all__ = ['AsyncIOMotorClient','AsyncIOMotorClientEncryption']
 
 
 def create_asyncio_class(cls):
@@ -70,3 +70,7 @@ def create_asyncio_class(cls):
 
 AsyncIOMotorGridOutCursor = create_asyncio_class(
     motor_gridfs.AgnosticGridOutCursor)
+
+
+AsyncIOMotorClientEncryption = create_asyncio_class(
+    core.AgnosticClientEncryption)

motor/motor_tornado.py

@@ -18,7 +18,7 @@
 from .frameworks import tornado as tornado_framework
 from .metaprogramming import create_class_with_framework
 
-__all__ = ['MotorClient']
+__all__ = ['MotorClient', 'MotorClientEncryption']
 
 
 def create_motor_class(cls):
@@ -60,3 +60,6 @@ def create_motor_class(cls):
 
 
 MotorGridOutCursor = create_motor_class(motor_gridfs.AgnosticGridOutCursor)
+
+
+MotorClientEncryption = create_motor_class(core.AgnosticClientEncryption)

setup.py

@@ -33,6 +33,10 @@
 
 tests_require = ['mockupdb>=1.4.0']
 
+extras_require = {
+    'encryption': ['pymongo[encryption]>=3.11,<4'],
+}
+
 class test(Command):
     description = "run the tests"
 
@@ -138,6 +142,7 @@ def run(self):
       url='https://github.com/mongodb/motor/',
       python_requires='>=3.5.2',
       install_requires=install_requires,
+      extras_require=extras_require,
       license='http://www.apache.org/licenses/LICENSE-2.0',
       classifiers=[c for c in classifiers.split('\n') if c],
       keywords=[

test/asyncio_tests/test_asyncio_encryption.py

@@ -0,0 +1,205 @@
+# Copyright 2021-present MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Test Explicit Encryption with AsyncIOMotorClient."""
+
+import unittest
+import uuid
+from bson.binary import (Binary,
+                         JAVA_LEGACY,
+                         STANDARD,
+                         UUID_SUBTYPE)
+from bson.codec_options import CodecOptions
+from bson.errors import BSONError
+
+from motor.motor_asyncio import AsyncIOMotorClientEncryption
+
+from pymongo.encryption import Algorithm
+from pymongo.errors import InvalidOperation
+from test.asyncio_tests import (asyncio_test,
+                                AsyncIOTestCase,
+                                skip_if_mongos)
+
+KMS_PROVIDERS = {'local': {'key': b'\x00'*96}}
+
+OPTS = CodecOptions(uuid_representation=STANDARD)
+
+try:
+    import pymongocrypt
+    _HAVE_PYMONGOCRYPT = True
+except ImportError:
+    _HAVE_PYMONGOCRYPT = False
+
+
+@unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is a required dependency")
+class TestExplicitSimple(AsyncIOTestCase):
+    def assertEncrypted(self, val):
+        self.assertIsInstance(val, Binary)
+        self.assertEqual(val.subtype, 6)
+
+    def assertBinaryUUID(self, val):
+        self.assertIsInstance(val, Binary)
+        self.assertEqual(val.subtype, UUID_SUBTYPE)
+
+    @asyncio_test
+    async def test_encrypt_decrypt(self):
+        client = self.asyncio_client()
+        client_encryption = AsyncIOMotorClientEncryption(
+            KMS_PROVIDERS, 'keyvault.datakeys', client, OPTS)
+        # Use standard UUID representation.
+        key_vault = client.keyvault.get_collection(
+            'datakeys', codec_options=OPTS)
+
+        # Create the encrypted field's data key.
+        key_id = await client_encryption.create_data_key(
+            'local', key_alt_names=['name'])
+        self.assertBinaryUUID(key_id)
+        self.assertTrue(await key_vault.find_one({'_id': key_id}))
+
+        # Create an unused data key to make sure filtering works.
+        unused_key_id = await client_encryption.create_data_key(
+            'local', key_alt_names=['unused'])
+        self.assertBinaryUUID(unused_key_id)
+        self.assertTrue(await key_vault.find_one({'_id': unused_key_id}))
+
+        doc = {'_id': 0, 'ssn': '000'}
+        encrypted_ssn = await client_encryption.encrypt(
+            doc['ssn'], Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic,
+            key_id=key_id)
+
+        # Ensure encryption via key_alt_name for the same key produces the
+        # same output.
+        encrypted_ssn2 = await client_encryption.encrypt(
+            doc['ssn'], Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic,
+            key_alt_name='name')
+        self.assertEqual(encrypted_ssn, encrypted_ssn2)
+
+        # Test decryption.
+        decrypted_ssn = await client_encryption.decrypt(encrypted_ssn)
+        self.assertEqual(decrypted_ssn, doc['ssn'])
+
+        await key_vault.drop()
+        await client_encryption.close()
+
+    @asyncio_test
+    async def test_validation(self):
+        client = self.asyncio_client()
+        client_encryption = AsyncIOMotorClientEncryption(
+            KMS_PROVIDERS, 'keyvault.datakeys', client, OPTS)
+
+        msg = 'value to decrypt must be a bson.binary.Binary with subtype 6'
+        with self.assertRaisesRegex(TypeError, msg):
+            await client_encryption.decrypt('str')
+        with self.assertRaisesRegex(TypeError, msg):
+            await client_encryption.decrypt(Binary(b'123'))
+
+        msg = 'key_id must be a bson.binary.Binary with subtype 4'
+        algo = Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic
+        with self.assertRaisesRegex(TypeError, msg):
+            await client_encryption.encrypt('str', algo, key_id=uuid.uuid4())
+        with self.assertRaisesRegex(TypeError, msg):
+            await client_encryption.encrypt('str', algo, key_id=Binary(b'123'))
+
+        await client_encryption.close()
+
+    @asyncio_test
+    async def test_bson_errors(self):
+        client = self.asyncio_client()
+        client_encryption = AsyncIOMotorClientEncryption(
+            KMS_PROVIDERS, 'keyvault.datakeys', client, OPTS)
+
+        # Attempt to encrypt an unencodable object.
+        unencodable_value = object()
+        with self.assertRaises(BSONError):
+            await client_encryption.encrypt(
+                unencodable_value,
+                Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic,
+                key_id=Binary(uuid.uuid4().bytes, UUID_SUBTYPE))
+
+        await client_encryption.close()
+
+    @asyncio_test
+    async def test_codec_options(self):
+        client = self.asyncio_client()
+        with self.assertRaisesRegex(TypeError, 'codec_options must be'):
+            AsyncIOMotorClientEncryption(
+                KMS_PROVIDERS, 'keyvault.datakeys', client, None)
+
+        opts = CodecOptions(uuid_representation=JAVA_LEGACY)
+        client_encryption_legacy = AsyncIOMotorClientEncryption(
+            KMS_PROVIDERS, 'keyvault.datakeys', client, opts)
+        # self.addCleanup(client_encryption_legacy.close)
+
+        # Create the encrypted field's data key.
+        key_id = await client_encryption_legacy.create_data_key('local')
+
+        # Encrypt a UUID with JAVA_LEGACY codec options.
+        value = uuid.uuid4()
+        encrypted_legacy = await client_encryption_legacy.encrypt(
+            value, Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic,
+            key_id=key_id)
+        decrypted_value_legacy = await client_encryption_legacy.decrypt(
+            encrypted_legacy)
+        self.assertEqual(decrypted_value_legacy, value)
+
+        # Encrypt the same UUID with STANDARD codec options.
+        client_encryption = AsyncIOMotorClientEncryption(
+            KMS_PROVIDERS, 'keyvault.datakeys', client, OPTS)
+        encrypted_standard = await client_encryption.encrypt(
+            value, Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic,
+            key_id=key_id)
+        decrypted_standard = await client_encryption.decrypt(encrypted_standard)
+        self.assertEqual(decrypted_standard, value)
+
+        # Test that codec_options is applied during encryption.
+        self.assertNotEqual(encrypted_standard, encrypted_legacy)
+        # Test that codec_options is applied during decryption.
+        self.assertEqual(
+            await client_encryption_legacy.decrypt(encrypted_standard), value)
+        self.assertNotEqual(
+            await client_encryption.decrypt(encrypted_legacy), value)
+
+        await client_encryption.close()
+
+    @asyncio_test
+    async def test_close(self):
+        client = self.asyncio_client()
+        client_encryption = AsyncIOMotorClientEncryption(
+            KMS_PROVIDERS, 'keyvault.datakeys', client, OPTS)
+        await client_encryption.close()
+        # Close can be called multiple times.
+        await client_encryption.close()
+        algo = Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic
+        msg = 'Cannot use closed ClientEncryption'
+        with self.assertRaisesRegex(InvalidOperation, msg):
+            await client_encryption.create_data_key('local')
+        with self.assertRaisesRegex(InvalidOperation, msg):
+            await client_encryption.encrypt('val', algo, key_alt_name='name')
+        with self.assertRaisesRegex(InvalidOperation, msg):
+            await client_encryption.decrypt(Binary(b'', 6))
+
+    @asyncio_test
+    async def test_with_statement(self):
+        client = self.asyncio_client()
+        async with AsyncIOMotorClientEncryption(
+                KMS_PROVIDERS, 'keyvault.datakeys',
+                client, OPTS) as client_encryption:
+            pass
+        with self.assertRaisesRegex(
+                InvalidOperation, 'Cannot use closed ClientEncryption'):
+            await client_encryption.create_data_key('local')
+
+
+if __name__ == '__main__':
+    unittest.main()