mongodb/mongo-python-driver
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

PYTHON-2267: Allow UUID key_id to be passed to ClientEncryption.encrypt (#1494)

Browse Source
This commit is contained in:
ilukyanchikov 2024-02-02 21:29:14 +03:00 committed by GitHub
parent c2af3dfeed
commit da2bf9dbcc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 21 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
pymongo/encryption.py
View File

@ -18,6 +18,7 @@ from __future__ import annotations
import contextlib
import enum
import socket
import uuid
import weakref
from copy import deepcopy
from typing import (
@ -30,6 +31,7 @@ from typing import (
MutableMapping,
Optional,
Sequence,
Union,
cast,
)
@ -759,7 +761,7 @@ class ClientEncryption(Generic[_DocumentType]):
self,
value: Any,
algorithm: str,
key_id: Optional[Binary] = None,
key_id: Optional[Union[Binary, uuid.UUID]] = None,
key_alt_name: Optional[str] = None,
query_type: Optional[str] = None,
contention_factor: Optional[int] = None,
@ -767,6 +769,8 @@ class ClientEncryption(Generic[_DocumentType]):
is_expression: bool = False,
) -> Any:
self._check_closed()
if isinstance(key_id, uuid.UUID):
key_id = Binary.from_uuid(key_id)
if key_id is not None and not (
isinstance(key_id, Binary) and key_id.subtype == UUID_SUBTYPE
):
@ -799,7 +803,7 @@ class ClientEncryption(Generic[_DocumentType]):
self,
value: Any,
algorithm: str,
key_id: Optional[Binary] = None,
key_id: Optional[Union[Binary, uuid.UUID]] = None,
key_alt_name: Optional[str] = None,
query_type: Optional[str] = None,
contention_factor: Optional[int] = None,
@ -826,6 +830,9 @@ class ClientEncryption(Generic[_DocumentType]):
:return: The encrypted value, a :class:`~bson.binary.Binary` with subtype 6.
.. versionchanged:: 4.7
``key_id`` can now be passed in as a :class:`uuid.UUID`.
.. versionchanged:: 4.2
Added the `query_type` and `contention_factor` parameters.
"""
@ -847,7 +854,7 @@ class ClientEncryption(Generic[_DocumentType]):
self,
expression: Mapping[str, Any],
algorithm: str,
key_id: Optional[Binary] = None,
key_id: Optional[Union[Binary, uuid.UUID]] = None,
key_alt_name: Optional[str] = None,
query_type: Optional[str] = None,
contention_factor: Optional[int] = None,
@ -875,6 +882,9 @@ class ClientEncryption(Generic[_DocumentType]):
:return: The encrypted expression, a :class:`~bson.RawBSONDocument`.
.. versionchanged:: 4.7
``key_id`` can now be passed in as a :class:`uuid.UUID`.
.. versionadded:: 4.4
"""
return cast(

11
test/test_encryption.py
View File

@ -461,6 +461,14 @@ class TestExplicitSimple(EncryptionIntegrationTest):
)
self.assertEqual(encrypted_ssn, encrypted_ssn2)
# Test encryption via UUID
encrypted_ssn3 = client_encryption.encrypt(
doc["ssn"],
Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic,
key_id=key_id.as_uuid(),
)
self.assertEqual(encrypted_ssn, encrypted_ssn3)
# Test decryption.
decrypted_ssn = client_encryption.decrypt(encrypted_ssn)
self.assertEqual(decrypted_ssn, doc["ssn"])
@ -479,9 +487,6 @@ class TestExplicitSimple(EncryptionIntegrationTest):
msg = "key_id must be a bson.binary.Binary with subtype 4"
algo = Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Deterministic
uid = uuid.uuid4()
with self.assertRaisesRegex(TypeError, msg):
client_encryption.encrypt("str", algo, key_id=uid) # type: ignore[arg-type]
with self.assertRaisesRegex(TypeError, msg):
client_encryption.encrypt("str", algo, key_id=Binary(b"123"))