PYTHON-3053 Key Management API

.evergreen/run-tests.sh

@@ -141,7 +141,8 @@ if [ -n "$TEST_ENCRYPTION" ]; then
  # TODO: Test with 'pip install pymongocrypt'
  git clone https://github.com/mongodb/libmongocrypt.git libmongocrypt_git
  python -m pip install --prefer-binary -r .evergreen/test-encryption-requirements.txt
- python -m pip install ./libmongocrypt_git/bindings/python
+ #python -m pip install ./libmongocrypt_git/bindings/python
+ python -m pip install pymongocrypt@git+ssh://git@github.com/blink1073/libmongocrypt.git@PYTHON-3053#subdirectory=bindings/python
  python -c "import pymongocrypt; print('pymongocrypt version: '+pymongocrypt.__version__)"
  python -c "import pymongocrypt; print('libmongocrypt version: '+pymongocrypt.libmongocrypt_version())"
  # PATH is updated by PREPARE_SHELL for access to mongocryptd.

pymongo/encryption.py

@@ -40,6 +40,7 @@
 from bson.raw_bson import DEFAULT_RAW_BSON_OPTIONS, RawBSONDocument, _inflate_bson
 from bson.son import SON
 from pymongo import _csot
+from pymongo.cursor import Cursor
 from pymongo.daemon import _spawn_daemon
 from pymongo.encryption_options import AutoEncryptionOpts
 from pymongo.errors import (
@@ -50,8 +51,10 @@
 )
 from pymongo.mongo_client import MongoClient
 from pymongo.network import BLOCKING_IO_ERRORS
+from pymongo.operations import UpdateOne
 from pymongo.pool import PoolOptions, _configured_socket
 from pymongo.read_concern import ReadConcern
+from pymongo.results import BulkWriteResult, DeleteResult
 from pymongo.ssl_support import get_ssl_context
 from pymongo.uri_parser import parse_host
 from pymongo.write_concern import WriteConcern
@@ -60,6 +63,7 @@
 _KMS_CONNECT_TIMEOUT = 10 # TODO: CDRIVER-3262 will define this value.
 _MONGOCRYPTD_TIMEOUT_MS = 10000
 
+
 _DATA_KEY_OPTS: CodecOptions = CodecOptions(document_class=SON, uuid_representation=STANDARD)
 # Use RawBSONDocument codec options to avoid needlessly decoding
 # documents from the key vault.
@@ -256,6 +260,21 @@ def close(self):
  self.mongocryptd_client = None
 
 
+class RewrapManyDataKeyResult(object):
+ def __init__(self, bulk_write_result: Optional[BulkWriteResult] = None) -> None:
+ """Result object returned by a ``rewrap_many_data_key`` operation.
+
+ :Parameters:
+ - `bulk_write_result`: he result of the bulk write operation used to update the key vault collection with one or more rewrapped data keys. If ``rewrap_many_data_key()`` does not find any matching keys to rewrap, no bulk write operation will be executed and this field will be ``None``.
+ """
+ self._bulk_write_result = bulk_write_result
+
+ @property
+ def bulk_write_result(self) -> Optional[BulkWriteResult]:
+ """The result of the bulk write operation used to update the key vault collection with one or more rewrapped data keys. If ``rewrap_many_data_key()`` does not find any matching keys to rewrap, no bulk write operation will be executed and this field will be ``None``."""
+ return self._bulk_write_result
+
+
 class _Encrypter(object):
  """Encrypts and decrypts MongoDB commands.
 
@@ -514,12 +533,18 @@ def __init__(
  self._encryption = ExplicitEncrypter(
  self._io_callbacks, MongoCryptOptions(kms_providers, None)
  )
+ # Create a version of the key vault collection that returns ObjectId
+ # objects instead of UUIDs.
+ self._key_vault_coll = self._io_callbacks.key_vault_coll.with_options(
+ codec_options=DEFAULT_RAW_BSON_OPTIONS,
+ )
 
  def create_data_key(
  self,
  kms_provider: str,
  master_key: Optional[Mapping[str, Any]] = None,
  key_alt_names: Optional[Sequence[str]] = None,
+ key_material: Optional[bytes] = None,
  ) -> Binary:
  """Create and insert a new data key into the key vault collection.
 
@@ -580,16 +605,24 @@ def create_data_key(
  # reference the key with the alternate name
  client_encryption.encrypt("457-55-5462", keyAltName="name1",
  algorithm=Algorithm.AEAD_AES_256_CBC_HMAC_SHA_512_Random)
+ - `key_material` (optional): Sets the custom key material to be used
+ by the data key for encryption and decryption.
 
  :Returns:
  The ``_id`` of the created data key document as a
  :class:`~bson.binary.Binary` with subtype
  :data:`~bson.binary.UUID_SUBTYPE`.
+
+ .. versionchanged:: 4.2
+ Added the `key_material` parameter.
  """
  self._check_closed()
  with _wrap_encryption_errors():
  return self._encryption.create_data_key(
- kms_provider, master_key=master_key, key_alt_names=key_alt_names
+ kms_provider,
+ master_key=master_key,
+ key_alt_names=key_alt_names,
+ key_material=key_material,
  )
 
  def encrypt(
@@ -674,6 +707,145 @@ def decrypt(self, value: Binary) -> Any:
  decrypted_doc = self._encryption.decrypt(doc)
  return decode(decrypted_doc, codec_options=self._codec_options)["v"]
 
+ def get_key(self, id: Binary) -> Optional[RawBSONDocument]:
+ """Get a data key by id.
+
+ :Parameters:
+ - `id` (Binary): The UUID of a key a which must be a
+ :class:`~bson.binary.Binary` with subtype 4 (
+ :attr:`~bson.binary.UUID_SUBTYPE`).
+
+ :Returns:
+ The key document.
+ """
+ self._check_closed()
+ return self._key_vault_coll.find_one({"_id": id})
+
+ def get_keys(self) -> Cursor[RawBSONDocument]:
+ """Get all of the data keys.
+
+ :Returns:
+ An instance of :class:`~pymongo.cursor.Cursor` over the data key
+ documents.
+ """
+ self._check_closed()
+ return self._key_vault_coll.find({})
+
+ def delete_key(self, id: Binary) -> DeleteResult:
+ """Delete a key document in the key vault collection that has the given ``key_id``.
+
+ :Parameters:
+ - `id` (Binary): The UUID of a key a which must be a
+ :class:`~bson.binary.Binary` with subtype 4 (
+ :attr:`~bson.binary.UUID_SUBTYPE`).
+
+ :Returns:
+ The delete result.
+ """
+ self._check_closed()
+ return self._key_vault_coll.delete_one({"_id": id})
+
+ def add_key_alt_name(self, id: Binary, key_alt_name: str) -> Any:
+ """Add ``key_alt_name`` to the set of alternate names in the key document with UUID ``key_id``.
+
+ :Parameters:
+ - ``id``: The UUID of a key a which must be a
+ :class:`~bson.binary.Binary` with subtype 4 (
+ :attr:`~bson.binary.UUID_SUBTYPE`).
+ - ``key_alt_name``: The key alternate name to add.
+
+ :Returns:
+ The previous version of the key document.
+ """
+ self._check_closed()
+ update = {"$addToSet": {"keyAltNames": key_alt_name}}
+ return self._key_vault_coll.find_one_and_update({"_id": id}, update)
+
+ def get_key_by_alt_name(self, key_alt_name: str) -> Optional[RawBSONDocument]:
+ """Get a key document in the key vault collection that has the given ``key_alt_name``.
+
+ :Parameters:
+ - `key_alt_name`: (str): The key alternate name of the key to get.
+
+ :Returns:
+ The key document.
+ """
+ self._check_closed()
+ return self._key_vault_coll.find_one({"keyAltNames": key_alt_name})
+
+ def remove_key_alt_name(self, id: Binary, key_alt_name: str) -> Optional[RawBSONDocument]:
+ """Remove ``key_alt_name`` from the set of keyAltNames in the key document with UUID ``id``.
+
+ Also removes the ``keyAltNames`` field from the key document if it would otherwise be empty.
+
+ :Parameters:
+ - ``id``: The UUID of a key a which must be a
+ :class:`~bson.binary.Binary` with subtype 4 (
+ :attr:`~bson.binary.UUID_SUBTYPE`).
+ - ``key_alt_name``: The key alternate name to remove.
+
+ :Returns:
+ Returns the previous version of the key document.
+ """
+ self._check_closed()
+ pipeline = [
+ {
+ "$set": {
+ "keyAltNames": {
+ "$cond": [
+ {"$eq": ["$keyAltNames", [key_alt_name]]},
+ "$$REMOVE",
+ {
+ "$filter": {
+ "input": "$keyAltNames",
+ "cond": {"$ne": ["$$this", key_alt_name]},
+ }
+ },
+ ]
+ }
+ }
+ }
+ ]
+ return self._key_vault_coll.find_one_and_update({"_id": id}, pipeline)
+
+ def rewrap_many_data_key(
+ self,
+ filter: Mapping[str, Any],
+ provider: Optional[str] = None,
+ master_key: Optional[Mapping[str, Any]] = None,
+ ) -> RewrapManyDataKeyResult:
+ """Decrypts and encrypts all matching data keys in the key vault with a possibly new `master_key` value.
+
+ :Parameters:
+ - `filter`: A document used to filter the data keys.
+ - `provider`: The new KMS provider to use to encrypt the data keys,
+ or ``None`` to use the current KMS provider(s).
+ - ``master_key``: The master key fields corresponding to the new KMS
+ provider when ``provider`` is not ``None``.
+
+ :Returns:
+ A :class:`RewrapManyDataKeyResult`.
+ """
+ self._check_closed()
+ with _wrap_encryption_errors():
+ raw_result = self._encryption.rewrap_many_data_key(filter, provider, master_key)
+ if raw_result is None:
+ return RewrapManyDataKeyResult()
+
+ raw_doc = RawBSONDocument(raw_result, DEFAULT_RAW_BSON_OPTIONS)
+ replacements = []
+ for key in raw_doc["v"]:
+ update_model = {
+ "$set": {"keyMaterial": key["keyMaterial"], "masterKey": key["masterKey"]},
+ "$currentDate": {"updateDate": True},
+ }
+ op = UpdateOne({"_id": key["_id"]}, update_model)
+ replacements.append(op)
+ if not replacements:
+ return RewrapManyDataKeyResult()
+ result = self._key_vault_coll.bulk_write(replacements)
+ return RewrapManyDataKeyResult(result)
+
  def __enter__(self) -> "ClientEncryption":
  return self
 

setup.py

@@ -277,7 +277,7 @@ def build_extension(self, ext):
 
 extras_require = {
  "encryption": [
- "pymongocrypt@git+ssh://git@github.com/mongodb/libmongocrypt.git@pymongocrypt-1.3.0b0#subdirectory=bindings/python"
+ "pymongocrypt@git+ssh://git@github.com/blink1073/libmongocrypt.git@PYTHON-3053#subdirectory=bindings/python"
  ],
  "ocsp": pyopenssl_reqs,
  "snappy": ["python-snappy"],

test/__init__.py

@@ -15,6 +15,7 @@
 """Test suite for pymongo, bson, and gridfs.
 """
 
+import base64
 import gc
 import os
 import socket
@@ -116,6 +117,27 @@
  COMPRESSORS = COMPRESSORS or "zlib"
 
 
+# Shared KMS data.
+LOCAL_MASTER_KEY = base64.b64decode(
+ b"Mng0NCt4ZHVUYUJCa1kxNkVyNUR1QURhZ2h2UzR2d2RrZzh0cFBwM3R6NmdWMDFBMUN3YkQ"
+ b"5aXRRMkhGRGdQV09wOGVNYUMxT2k3NjZKelhaQmRCZGJkTXVyZG9uSjFk"
+)
+AWS_CREDS = {
+ "accessKeyId": os.environ.get("FLE_AWS_KEY", ""),
+ "secretAccessKey": os.environ.get("FLE_AWS_SECRET", ""),
+}
+AZURE_CREDS = {
+ "tenantId": os.environ.get("FLE_AZURE_TENANTID", ""),
+ "clientId": os.environ.get("FLE_AZURE_CLIENTID", ""),
+ "clientSecret": os.environ.get("FLE_AZURE_CLIENTSECRET", ""),
+}
+GCP_CREDS = {
+ "email": os.environ.get("FLE_GCP_EMAIL", ""),
+ "privateKey": os.environ.get("FLE_GCP_PRIVATEKEY", ""),
+}
+KMIP_CREDS = {"endpoint": os.environ.get("FLE_KMIP_ENDPOINT", "localhost:5698")}
+
+
 def is_server_resolvable():
  """Returns True if 'server' is resolvable."""
  socket_timeout = socket.getdefaulttimeout()

test/client-side-encryption/spec/unified/addKeyAltName.json

@@ -22,7 +22,11 @@
  "keyVaultClient": "client0",
  "keyVaultNamespace": "keyvault.datakeys",
  "kmsProviders": {
- "local": {}
+ "local": {
+ "key": {
+ "$$placeholder": 1
+ }
+ }
  }
  }
  }

test/client-side-encryption/spec/unified/createKey-kms_providers-invalid.json → test/client-side-encryption/spec/unified/createDataKey-kms_providers-invalid.json

@@ -1,5 +1,5 @@
 {
- "description": "createKey-provider-invalid",
+ "description": "createDataKey-provider-invalid",
  "schemaVersion": "1.8",
  "runOnRequirements": [
  {
@@ -24,7 +24,14 @@
  "keyVaultClient": "client0",
  "keyVaultNamespace": "keyvault.datakeys",
  "kmsProviders": {
- "aws": {}
+ "aws": {
+ "accessKeyId": {
+ "$$placeholder": 1
+ },
+ "secretAccessKey": {
+ "$$placeholder": 1
+ }
+ }
  }
  }
  }
@@ -35,7 +42,7 @@
  "description": "create data key without required master key fields",
  "operations": [
  {
- "name": "createKey",
+ "name": "createDataKey",
  "object": "clientEncryption0",
  "arguments": {
  "kmsProvider": "aws",
@@ -59,7 +66,7 @@
  "description": "create data key with invalid master key field",
  "operations": [
  {
- "name": "createKey",
+ "name": "createDataKey",
  "object": "clientEncryption0",
  "arguments": {
  "kmsProvider": "local",
@@ -85,7 +92,7 @@
  "description": "create data key with invalid master key",
  "operations": [
  {
- "name": "createKey",
+ "name": "createDataKey",
  "object": "clientEncryption0",
  "arguments": {
  "kmsProvider": "aws",