PYTHON-1884 Support auto encryption in network.command

ShaneHarvey · ShaneHarvey · commit 8888e977f8ec · 2019-08-06T14:54:23.000-07:00
Add encryption I/O callbacks.
Add mongocryptd process management.
Add simple test for auto encryption round trip.
diff --git a/bson/raw_bson.py b/bson/raw_bson.py
@@ -95,8 +95,8 @@ def __inflated(self):
             # We already validated the object's size when this document was
             # created, so no need to do that again.
             # Use SON to preserve ordering of elements.
-            self.__inflated_doc = _raw_to_dict(
-                self.__raw, 4, len(self.__raw)-1, self.__codec_options, SON())
+            self.__inflated_doc = _inflate_bson(
+                self.__raw, self.__codec_options)
         return self.__inflated_doc
 
     def __getitem__(self, item):
@@ -118,6 +118,20 @@ def __repr__(self):
                 % (self.raw, self.__codec_options))
 
 
+def _inflate_bson(bson_bytes, codec_options):
+    """Inflates the top level fields of a BSON document.
+
+    :Parameters:
+      - `bson_bytes`: the BSON bytes that compose this document
+      - `codec_options`: An instance of
+        :class:`~bson.codec_options.CodecOptions` whose ``document_class``
+        must be :class:`RawBSONDocument`.
+    """
+    # Use SON to preserve ordering of elements.
+    return _raw_to_dict(
+        bson_bytes, 4, len(bson_bytes)-1, codec_options, SON())
+
+
 DEFAULT_RAW_BSON_OPTIONS = DEFAULT.with_options(document_class=RawBSONDocument)
 """The default :class:`~bson.codec_options.CodecOptions` for
 :class:`RawBSONDocument`.
diff --git a/pymongo/encryption.py b/pymongo/encryption.py
@@ -0,0 +1,250 @@
+# Copyright 2019-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.
+
+"""Client side encryption implementation."""
+
+import subprocess
+import weakref
+
+from pymongocrypt.auto_encrypter import AutoEncrypter
+from pymongocrypt.errors import MongoCryptError
+from pymongocrypt.mongocrypt import MongoCryptOptions
+from pymongocrypt.state_machine import MongoCryptCallback
+
+from bson import _bson_to_dict, _dict_to_bson
+from bson.binary import STANDARD
+from bson.codec_options import CodecOptions
+from bson.raw_bson import (DEFAULT_RAW_BSON_OPTIONS,
+                           RawBSONDocument,
+                           _inflate_bson)
+from bson.son import SON
+
+from pymongo.errors import ServerSelectionTimeoutError
+from pymongo.mongo_client import MongoClient
+from pymongo.pool import _configured_socket, PoolOptions
+from pymongo.ssl_support import get_ssl_context
+
+
+_HTTPS_PORT = 443
+_KMS_CONNECT_TIMEOUT = 10  # TODO: CDRIVER-3262 will define this value.
+_MONGOCRYPTD_TIMEOUT_MS = 1000
+
+_DATA_KEY_OPTS = CodecOptions(document_class=SON, uuid_representation=STANDARD)
+# Use RawBSONDocument codec options to avoid needlessly decoding
+# documents from the key vault.
+_KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument,
+                               uuid_representation=STANDARD)
+
+
+class _EncryptionIO(MongoCryptCallback):
+    def __init__(self, client, key_vault_coll, mongocryptd_client, opts):
+        """Internal class to perform I/O on behalf of pymongocrypt."""
+        # Use a weak ref to break reference cycle.
+        self.client_ref = weakref.ref(client)
+        self.key_vault_coll = key_vault_coll.with_options(
+            codec_options=_KEY_VAULT_OPTS)
+        self.mongocryptd_client = mongocryptd_client
+        self.opts = opts
+        self._spawned = False
+
+    def kms_request(self, kms_context):
+        """Complete a KMS request.
+
+        :Parameters:
+          - `kms_context`: A :class:`MongoCryptKmsContext`.
+
+        :Returns:
+          None
+        """
+        endpoint = kms_context.endpoint
+        message = kms_context.message
+        ctx = get_ssl_context(None, None, None, None, None, None, True)
+        opts = PoolOptions(connect_timeout=_KMS_CONNECT_TIMEOUT,
+                           socket_timeout=_KMS_CONNECT_TIMEOUT,
+                           ssl_context=ctx)
+        try:
+            conn = _configured_socket((endpoint, _HTTPS_PORT), opts)
+            conn.sendall(message)
+        except Exception as exc:
+            raise MongoCryptError(str(exc))
+
+        while kms_context.bytes_needed > 0:
+            data = conn.recv(kms_context.bytes_needed)
+            kms_context.feed(data)
+
+    def collection_info(self, database, filter):
+        """Get the collection info for a namespace.
+
+        The returned collection info is passed to libmongocrypt which reads
+        the JSON schema.
+
+        :Parameters:
+          - `database`: The database on which to run listCollections.
+          - `filter`: The filter to pass to listCollections.
+
+        :Returns:
+          The first document from the listCollections command response as BSON.
+        """
+        with self.client_ref()[database].list_collections(
+                filter=RawBSONDocument(filter)) as cursor:
+            for doc in cursor:
+                return _dict_to_bson(doc, False, _DATA_KEY_OPTS)
+
+    def spawn(self):
+        """Spawn mongocryptd.
+
+        Note this method is thread safe; at most one mongocryptd will start
+        successfully.
+        """
+        self._spawned = True
+        args = [self.opts._mongocryptd_spawn_path or 'mongocryptd']
+        args.extend(self.opts._mongocryptd_spawn_args)
+        subprocess.Popen(args)
+
+    def mark_command(self, database, cmd):
+        """Mark a command for encryption.
+
+        :Parameters:
+          - `database`: The database on which to run this command.
+          - `cmd`: The BSON command to run.
+
+        :Returns:
+          The marked command response from mongocryptd.
+        """
+        if not self._spawned and not self.opts._mongocryptd_bypass_spawn:
+            self.spawn()
+        # Database.command only supports mutable mappings so we need to decode
+        # the raw BSON command first.
+        inflated_cmd = _inflate_bson(cmd, DEFAULT_RAW_BSON_OPTIONS)
+        try:
+            res = self.mongocryptd_client[database].command(
+                inflated_cmd,
+                codec_options=DEFAULT_RAW_BSON_OPTIONS)
+        except ServerSelectionTimeoutError:
+            if self.opts._mongocryptd_bypass_spawn:
+                raise
+            self.spawn()
+            res = self.mongocryptd_client[database].command(
+                inflated_cmd,
+                codec_options=DEFAULT_RAW_BSON_OPTIONS)
+        return res.raw
+
+    def fetch_keys(self, filter):
+        """Yields one or more keys from the key vault.
+
+        :Parameters:
+          - `filter`: The filter to pass to find.
+
+        :Returns:
+          A generator which yields the requested keys from the key vault.
+        """
+        with self.key_vault_coll.find(RawBSONDocument(filter)) as cursor:
+            for key in cursor:
+                yield key.raw
+
+    def insert_data_key(self, data_key):
+        """Insert a data key into the key vault.
+
+        :Parameters:
+          - `data_key`: The data key document to insert.
+
+        :Returns:
+          The _id of the inserted data key document.
+        """
+        # insert does not return the inserted _id when given a RawBSONDocument.
+        doc = _bson_to_dict(data_key, _DATA_KEY_OPTS)
+        res = self.key_vault_coll.insert_one(doc)
+        return res.inserted_id
+
+    def close(self):
+        """Release resources.
+
+        Note it is not safe to call this method from __del__ or any GC hooks.
+        """
+        self.client_ref = None
+        self.key_vault_coll = None
+        self.mongocryptd_client.close()
+        self.mongocryptd_client = None
+
+
+class _Encrypter(object):
+    def __init__(self, io_callbacks, opts):
+        """Encrypts and decrypts MongoDB commands.
+
+        This class is used to support automatic encryption and decryption of
+        MongoDB commands.
+
+        :Parameters:
+          - `io_callbacks`: A :class:`MongoCryptCallback`.
+          - `opts`: The encrypted client's :class:`AutoEncryptionOpts`.
+        """
+        if opts._schema_map is None:
+            schema_map = None
+        else:
+            schema_map = _dict_to_bson(opts._schema_map, False, _DATA_KEY_OPTS)
+        self._auto_encrypter = AutoEncrypter(io_callbacks, MongoCryptOptions(
+            opts._kms_providers, schema_map))
+        self._bypass_auto_encryption = opts._bypass_auto_encryption
+
+    def encrypt(self, database, cmd):
+        """Encrypt a MongoDB command.
+
+        :Parameters:
+          - `database`: The database for this command.
+          - `cmd`: A command as BSON.
+
+        :Returns:
+          The encrypted command to execute.
+        """
+        encrypted_cmd = self._auto_encrypter.encrypt(database, cmd)
+        # TODO: PYTHON-1922 avoid decoding the encrypted_cmd.
+        return _inflate_bson(encrypted_cmd, DEFAULT_RAW_BSON_OPTIONS)
+
+    def decrypt(self, response):
+        """Decrypt a MongoDB command response.
+
+        :Parameters:
+          - `response`: A MongoDB command response as BSON.
+
+        :Returns:
+          The decrypted command response.
+        """
+        return self._auto_encrypter.decrypt(response)
+
+    def close(self):
+        """Cleanup resources."""
+        self._auto_encrypter.close()
+
+    @staticmethod
+    def create(client, opts):
+        """Create a _CommandEncyptor for a client.
+
+        :Parameters:
+          - `client`: The encrypted MongoClient.
+          - `opts`: The encrypted client's :class:`AutoEncryptionOpts`.
+
+        :Returns:
+          A :class:`_CommandEncrypter` for this client.
+        """
+        key_vault_client = opts._key_vault_client or client
+        db, coll = opts._key_vault_namespace.split('.', 1)
+        key_vault_coll = key_vault_client[db][coll]
+
+        mongocryptd_client = MongoClient(
+            opts._mongocryptd_uri, connect=False,
+            serverSelectionTimeoutMS=_MONGOCRYPTD_TIMEOUT_MS)
+
+        io_callbacks = _EncryptionIO(
+            client, key_vault_coll, mongocryptd_client, opts)
+        return _Encrypter(io_callbacks, opts)
diff --git a/pymongo/message.py b/pymongo/message.py
@@ -1439,6 +1439,11 @@ def command_response(self):
         assert self.number_returned == 1
         return docs[0]
 
+    def raw_command_response(self):
+        """Return the bytes of the command response."""
+        # This should never be called on _OpReply.
+        raise NotImplementedError
+
     @classmethod
     def unpack(cls, msg):
         """Construct an _OpReply from raw bytes."""
@@ -1485,6 +1490,10 @@ def command_response(self):
         """Unpack a command response."""
         return self.unpack_response()[0]
 
+    def raw_command_response(self):
+        """Return the bytes of the command response."""
+        return self.payload_document
+
     @classmethod
     def unpack(cls, msg):
         """Construct an _OpMsg from raw bytes."""
diff --git a/pymongo/mongo_client.py b/pymongo/mongo_client.py
@@ -713,6 +713,12 @@ def target():
         self._kill_cursors_executor = executor
         executor.open()
 
+        self._encrypter = None
+        if self.__options.auto_encryption_opts:
+            from pymongo.encryption import _Encrypter
+            self._encrypter = _Encrypter.create(
+                self, self.__options.auto_encryption_opts)
+
     def _cache_credentials(self, source, credentials, connect=False):
         """Save a set of authentication credentials.
 
@@ -1152,6 +1158,9 @@ def close(self):
         self._kill_cursors_executor.close()
         self._process_periodic_tasks()
         self._topology.close()
+        if self._encrypter:
+            # TODO: PYTHON-1921 Encrypted MongoClients cannot be re-opened.
+            self._encrypter.close()
 
     def set_cursor_manager(self, manager_class):
         """DEPRECATED - Set this client's cursor manager.
diff --git a/pymongo/network.py b/pymongo/network.py
@@ -34,6 +34,7 @@
 except ImportError:
     _SELECT_ERROR = OSError
 
+from bson import _dict_to_bson, _decode_all_selective
 from bson.py3compat import PY3
 
 from pymongo import helpers, message
@@ -114,6 +115,18 @@ def command(sock, dbname, spec, slave_ok, is_mongos,
     if compression_ctx and name.lower() in _NO_COMPRESSION:
         compression_ctx = None
 
+    if (client and client._encrypter and
+            not client._encrypter._bypass_auto_encryption):
+        # Workaround for $clusterTime which is incompatible with check_keys.
+        if check_keys:
+            cluster_time = spec.pop('$clusterTime', None)
+        spec = orig = client._encrypter.encrypt(
+            dbname, _dict_to_bson(spec, check_keys, codec_options))
+        if check_keys and cluster_time:
+            spec['$clusterTime'] = cluster_time
+        # We already checked the keys, no need to do it again.
+        check_keys = False
+
     if use_op_msg:
         flags = 2 if unacknowledged else 0
         request_id, msg, size, max_doc_size = message._op_msg(
@@ -143,6 +156,7 @@ def command(sock, dbname, spec, slave_ok, is_mongos,
         sock.sendall(msg)
         if use_op_msg and unacknowledged:
             # Unacknowledged, fake a successful command response.
+            reply = None
             response_doc = {"ok": 1}
         else:
             reply = receive_message(sock, request_id)
@@ -170,6 +184,12 @@ def command(sock, dbname, spec, slave_ok, is_mongos,
         duration = (datetime.datetime.now() - start) + encoding_duration
         listeners.publish_command_success(
             duration, response_doc, name, request_id, address)
+
+    if client and client._encrypter and reply:
+        decrypted = client._encrypter.decrypt(reply.raw_command_response())
+        response_doc = _decode_all_selective(decrypted, codec_options,
+                                             user_fields)[0]
+
     return response_doc
 
 _UNPACK_COMPRESSION_HEADER = struct.Struct("<iiB").unpack
diff --git a/test/client-side-encryption/key-document-local.json b/test/client-side-encryption/key-document-local.json
@@ -0,0 +1,18 @@
+{
+  "_id": {
+    "$binary": {
+        "base64": "YWFhYWFhYWFhYWFhYWFhYQ==",
+        "subType": "04"
+    }
+},
+  "keyMaterial": {
+    "$binary": {
+      "base64": "db27rshiqK4Jqhb2xnwK4RfdFb9JuKeUe6xt5aYQF4o62tS75b7B4wxVN499gND9UVLUbpVKoyUoaZAeA895OENP335b8n8OwchcTFqS44t+P3zmhteYUQLIWQXaIgon7gEgLeJbaDHmSXS6/7NbfDDFlB37N7BP/2hx1yCOTN6NG/8M1ppw3LYT3CfP6EfXVEttDYtPbJpbb7nBVlxD7w==",
+      "subType": "00"
+    }
+  },
+  "creationDate": { "$date": { "$numberLong": "1564354142963" } },
+  "updateDate": { "$date": { "$numberLong": "1564354142963" } },
+  "status": { "$numberInt": "0" },
+  "masterKey": { "provider": "local" }
+}
diff --git a/test/client-side-encryption/schema.json b/test/client-side-encryption/schema.json
diff --git a/test/test_encryption.py b/test/test_encryption.py
