PYTHON-1947 Add codec_options to ClientEncryption

Author: ShaneHarvey

pymongo/encryption.py

@@ -309,7 +309,8 @@ class Algorithm(object):
 class ClientEncryption(object):
     """Explicit client side encryption."""
 
-    def __init__(self, kms_providers, key_vault_namespace, key_vault_client):
+    def __init__(self, kms_providers, key_vault_namespace, key_vault_client,
+                 codec_options):
         """Explicit client side encryption.
 
         The ClientEncryption class encapsulates explicit operations on a key
@@ -340,6 +341,9 @@ def __init__(self, kms_providers, key_vault_namespace, key_vault_client):
             provider.
           - `key_vault_client`: A MongoClient connected to a MongoDB cluster
             containing the `key_vault_namespace` collection.
+          - `codec_options`: An instance of
+            :class:`~bson.codec_options.CodecOptions` to use when encoding a
+            value for encryption and decoding the decrypted BSON value.
 
         .. versionadded:: 3.9
         """
@@ -349,9 +353,14 @@ def __init__(self, kms_providers, key_vault_namespace, key_vault_client):
                 "install a compatible version with: "
                 "python -m pip install pymongo['encryption']")
 
+        if not isinstance(codec_options, CodecOptions):
+            raise TypeError("codec_options must be an instance of "
+                            "bson.codec_options.CodecOptions")
+
         self._kms_providers = kms_providers
         self._key_vault_namespace = key_vault_namespace
         self._key_vault_client = key_vault_client
+        self._codec_options = codec_options
 
         db, coll = key_vault_namespace.split('.', 1)
         key_vault_coll = key_vault_client[db][coll]
@@ -413,8 +422,7 @@ def encrypt(self, value, algorithm, key_id=None, key_alt_name=None):
         :Returns:
           The encrypted value, a :class:`~bson.binary.Binary` with subtype 6.
         """
-        # TODO: Add a required codec_options argument for encoding?
-        doc = encode({'v': value}, codec_options=_DATA_KEY_OPTS)
+        doc = encode({'v': value}, codec_options=self._codec_options)
         if isinstance(key_id, uuid.UUID):
             raw_key_id = key_id.bytes
         else:
@@ -428,8 +436,7 @@ def _decrypt(self, value):
         """Internal decrypt helper."""
         doc = encode({'v': value})
         decrypted_doc = self._encryption.decrypt(doc)
-        # TODO: Add a required codec_options argument for decoding?
-        return decode(decrypted_doc, codec_options=_DATA_KEY_OPTS)['v']
+        return decode(decrypted_doc, codec_options=self._codec_options)['v']
 
     def decrypt(self, value):
         """Decrypt an encrypted value.

test/test_encryption.py

@@ -25,7 +25,10 @@
 sys.path[0:0] = [""]
 
 from bson import BSON, json_util
-from bson.binary import STANDARD, Binary, UUID_SUBTYPE
+from bson.binary import (Binary,
+                         JAVA_LEGACY,
+                         STANDARD,
+                         UUID_SUBTYPE)
 from bson.codec_options import CodecOptions
 from bson.errors import BSONError
 from bson.json_util import JSONOptions
@@ -35,6 +38,8 @@
 from pymongo.errors import (ConfigurationError,
                             EncryptionError,
                             OperationFailure)
+from pymongo.encryption import (Algorithm,
+                                ClientEncryption)
 from pymongo.encryption_options import AutoEncryptionOpts, _HAVE_PYMONGOCRYPT
 from pymongo.mongo_client import MongoClient
 from pymongo.write_concern import WriteConcern
@@ -47,16 +52,6 @@
 from test.utils_spec_runner import SpecRunner
 
 
-if _HAVE_PYMONGOCRYPT:
-    # Load the mongocrypt library.
-    from pymongocrypt.binding import init
-    init(os.environ.get('MONGOCRYPT_LIB', 'mongocrypt'))
-
-# This has to be imported after calling init().
-from pymongo.encryption import (Algorithm,
-                                ClientEncryption)
-
-
 def get_client_opts(client):
     return client._MongoClient__options
 
@@ -261,7 +256,7 @@ class TestExplicitSimple(EncryptionIntegrationTest):
 
     def test_encrypt_decrypt(self):
         client_encryption = ClientEncryption(
-            KMS_PROVIDERS, 'admin.datakeys', client_context.client)
+            KMS_PROVIDERS, 'admin.datakeys', client_context.client, OPTS)
         self.addCleanup(client_encryption.close)
         # Use standard UUID representation.
         key_vault = client_context.client.admin.get_collection(
@@ -296,7 +291,7 @@ def test_encrypt_decrypt(self):
 
     def test_validation(self):
         client_encryption = ClientEncryption(
-            KMS_PROVIDERS, 'admin.datakeys', client_context.client)
+            KMS_PROVIDERS, 'admin.datakeys', client_context.client, OPTS)
         self.addCleanup(client_encryption.close)
 
         msg = 'value to decrypt must be a bson.binary.Binary with subtype 6'
@@ -307,7 +302,7 @@ def test_validation(self):
 
     def test_bson_errors(self):
         client_encryption = ClientEncryption(
-            KMS_PROVIDERS, 'admin.datakeys', client_context.client)
+            KMS_PROVIDERS, 'admin.datakeys', client_context.client, OPTS)
         self.addCleanup(client_encryption.close)
 
         # Attempt to encrypt an unencodable object.
@@ -317,6 +312,43 @@ def test_bson_errors(self):
                 unencodable_value, Algorithm.Deterministic,
                 key_id=Binary(uuid.uuid4().bytes, UUID_SUBTYPE))
 
+    def test_codec_options(self):
+        with self.assertRaisesRegex(TypeError, 'codec_options must be'):
+            ClientEncryption(
+                KMS_PROVIDERS, 'admin.datakeys', client_context.client, None)
+
+        opts = CodecOptions(uuid_representation=JAVA_LEGACY)
+        client_encryption_legacy = ClientEncryption(
+            KMS_PROVIDERS, 'admin.datakeys', client_context.client, opts)
+        self.addCleanup(client_encryption_legacy.close)
+
+        # Create the encrypted field's data key.
+        key_id = client_encryption_legacy.create_data_key('local')
+
+        # Encrypt a UUID with JAVA_LEGACY codec options.
+        value = uuid.uuid4()
+        encrypted_legacy = client_encryption_legacy.encrypt(
+            value, Algorithm.Deterministic, key_id=key_id)
+        decrypted_value_legacy = client_encryption_legacy.decrypt(
+            encrypted_legacy)
+        self.assertEqual(decrypted_value_legacy, value)
+
+        # Encrypt the same UUID with STANDARD codec options.
+        client_encryption = ClientEncryption(
+            KMS_PROVIDERS, 'admin.datakeys', client_context.client, OPTS)
+        self.addCleanup(client_encryption.close)
+        encrypted_standard = client_encryption.encrypt(
+            value, Algorithm.Deterministic, key_id=key_id)
+        decrypted_standard = 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(
+            client_encryption_legacy.decrypt(encrypted_standard), value)
+        self.assertNotEqual(
+            client_encryption.decrypt(encrypted_legacy), value)
 
 # Spec tests
 
@@ -485,7 +517,8 @@ def test_data_key(self):
         self.addCleanup(client_encrypted.close)
 
         client_encryption = ClientEncryption(
-            self.kms_providers(), 'admin.datakeys', client_context.client)
+            self.kms_providers(), 'admin.datakeys', client_context.client,
+            OPTS)
         self.addCleanup(client_encryption.close)
 
         # Local create data key.
@@ -577,7 +610,7 @@ def _test_external_key_vault(self, with_external_key_vault):
         self.addCleanup(client_encrypted.close)
 
         client_encryption = ClientEncryption(
-            self.kms_providers(), 'admin.datakeys', key_vault_client)
+            self.kms_providers(), 'admin.datakeys', key_vault_client, OPTS)
         self.addCleanup(client_encryption.close)
 
         if with_external_key_vault:
@@ -684,7 +717,8 @@ def _test_corpus(self, opts):
         self.addCleanup(client_encrypted.close)
 
         client_encryption = ClientEncryption(
-            self.kms_providers(), 'admin.datakeys', client_context.client)
+            self.kms_providers(), 'admin.datakeys', client_context.client,
+            OPTS)
         self.addCleanup(client_encryption.close)
 
         corpus = self.fix_up_curpus(json_data('corpus', 'corpus.json'))