PYTHON-3003 Add kms_tls_options to configure options for KMS provider connections (#784)

ShaneHarvey · web-flow · commit 370e1652ad97 · 2021-11-10T16:49:31.000-08:00
diff --git a/pymongo/common.py b/pymongo/common.py
@@ -30,7 +30,6 @@
                                          validate_zlib_compression_level)
 from pymongo.driver_info import DriverInfo
 from pymongo.server_api import ServerApi
-from pymongo.encryption_options import validate_auto_encryption_opts_or_none
 from pymongo.errors import ConfigurationError
 from pymongo.monitoring import _validate_event_listeners
 from pymongo.read_concern import ReadConcern
@@ -582,6 +581,18 @@ def validate_tzinfo(dummy, value):
     return value
 
 
+def validate_auto_encryption_opts_or_none(option, value):
+    """Validate the driver keyword arg."""
+    if value is None:
+        return value
+    from pymongo.encryption_options import AutoEncryptionOpts
+    if not isinstance(value, AutoEncryptionOpts):
+        raise TypeError("%s must be an instance of AutoEncryptionOpts" % (
+            option,))
+
+    return value
+
+
 # Dictionary where keys are the names of public URI options, and values
 # are lists of aliases for that option.
 URI_OPTIONS_ALIAS_MAP = {
diff --git a/pymongo/encryption.py b/pymongo/encryption.py
@@ -46,6 +46,7 @@
                             EncryptionError,
                             InvalidOperation,
                             ServerSelectionTimeoutError)
+from pymongo.encryption_options import AutoEncryptionOpts
 from pymongo.mongo_client import MongoClient
 from pymongo.pool import _configured_socket, PoolOptions
 from pymongo.read_concern import ReadConcern
@@ -106,20 +107,23 @@ def kms_request(self, kms_context):
         """
         endpoint = kms_context.endpoint
         message = kms_context.message
-        host, port = parse_host(endpoint, _HTTPS_PORT)
-        # Enable strict certificate verification, OCSP, match hostname, and
-        # SNI using the system default CA certificates.
-        ctx = get_ssl_context(
-            None,   # certfile
-            None,   # passphrase
-            None,   # ca_certs
-            None,   # crlfile
-            False,  # allow_invalid_certificates
-            False,  # allow_invalid_hostnames
-            False)  # disable_ocsp_endpoint_check
+        provider = kms_context.kms_provider
+        ctx = self.opts._kms_ssl_contexts.get(provider)
+        if not ctx:
+            # Enable strict certificate verification, OCSP, match hostname, and
+            # SNI using the system default CA certificates.
+            ctx = get_ssl_context(
+                None,   # certfile
+                None,   # passphrase
+                None,   # ca_certs
+                None,   # crlfile
+                False,  # allow_invalid_certificates
+                False,  # allow_invalid_hostnames
+                False)  # disable_ocsp_endpoint_check
         opts = PoolOptions(connect_timeout=_KMS_CONNECT_TIMEOUT,
                            socket_timeout=_KMS_CONNECT_TIMEOUT,
                            ssl_context=ctx)
+        host, port = parse_host(endpoint, _HTTPS_PORT)
         conn = _configured_socket((host, port), opts)
         try:
             conn.sendall(message)
@@ -359,7 +363,7 @@ class ClientEncryption(object):
     """Explicit client-side field level encryption."""
 
     def __init__(self, kms_providers, key_vault_namespace, key_vault_client,
-                 codec_options):
+                 codec_options, kms_tls_options=None):
         """Explicit client-side field level encryption.
 
         The ClientEncryption class encapsulates explicit operations on a key
@@ -411,6 +415,16 @@ def __init__(self, kms_providers, key_vault_namespace, key_vault_client,
             should be the same CodecOptions instance configured on the
             MongoClient, Database, or Collection used to access application
             data.
+          - `kms_tls_options` (optional): A map of KMS provider names to TLS
+            options to use when creating secure connections to KMS providers.
+            Accepts the same TLS options as
+            :class:`pymongo.mongo_client.MongoClient`. For example, to
+            override the system default CA file::
+
+              kms_tls_options={'kmip': {'tlsCAFile': certifi.where()}}
+
+        .. versionchanged:: 4.0
+           Added the `kms_tls_options` parameter.
 
         .. versionadded:: 3.9
         """
@@ -432,7 +446,9 @@ def __init__(self, kms_providers, key_vault_namespace, key_vault_client,
         db, coll = key_vault_namespace.split('.', 1)
         key_vault_coll = key_vault_client[db][coll]
 
-        self._io_callbacks = _EncryptionIO(None, key_vault_coll, None, None)
+        opts = AutoEncryptionOpts(kms_providers, key_vault_namespace,
+                                  kms_tls_options=kms_tls_options)
+        self._io_callbacks = _EncryptionIO(None, key_vault_coll, None, opts)
         self._encryption = ExplicitEncrypter(
             self._io_callbacks, MongoCryptOptions(kms_providers, None))
 
diff --git a/pymongo/encryption_options.py b/pymongo/encryption_options.py
@@ -23,6 +23,7 @@
     _HAVE_PYMONGOCRYPT = False
 
 from pymongo.errors import ConfigurationError
+from pymongo.uri_parser import _parse_kms_tls_options
 
 
 class AutoEncryptionOpts(object):
@@ -35,7 +36,8 @@ def __init__(self, kms_providers, key_vault_namespace,
                  mongocryptd_uri='mongodb://localhost:27020',
                  mongocryptd_bypass_spawn=False,
                  mongocryptd_spawn_path='mongocryptd',
-                 mongocryptd_spawn_args=None):
+                 mongocryptd_spawn_args=None,
+                 kms_tls_options=None):
         """Options to configure automatic client-side field level encryption.
 
         Automatic client-side field level encryption requires MongoDB 4.2
@@ -118,6 +120,16 @@ def __init__(self, kms_providers, key_vault_namespace,
             ``['--idleShutdownTimeoutSecs=60']``. If the list does not include
             the ``idleShutdownTimeoutSecs`` option then
             ``'--idleShutdownTimeoutSecs=60'`` will be added.
+          - `kms_tls_options` (optional):  A map of KMS provider names to TLS
+            options to use when creating secure connections to KMS providers.
+            Accepts the same TLS options as
+            :class:`pymongo.mongo_client.MongoClient`. For example, to
+            override the system default CA file::
+
+              kms_tls_options={'kmip': {'tlsCAFile': certifi.where()}}
+
+        .. versionchanged:: 4.0
+           Added the `kms_tls_options` parameter.
 
         .. versionadded:: 3.9
         """
@@ -142,14 +154,5 @@ def __init__(self, kms_providers, key_vault_namespace,
         if not any('idleShutdownTimeoutSecs' in s
                    for s in self._mongocryptd_spawn_args):
             self._mongocryptd_spawn_args.append('--idleShutdownTimeoutSecs=60')
-
-
-def validate_auto_encryption_opts_or_none(option, value):
-    """Validate the driver keyword arg."""
-    if value is None:
-        return value
-    if not isinstance(value, AutoEncryptionOpts):
-        raise TypeError("%s must be an instance of AutoEncryptionOpts" % (
-            option,))
-
-    return value
+        # Maps KMS provider name to a SSLContext.
+        self._kms_ssl_contexts = _parse_kms_tls_options(kms_tls_options)
diff --git a/pymongo/uri_parser.py b/pymongo/uri_parser.py
@@ -20,6 +20,7 @@
 
 from urllib.parse import unquote_plus
 
+from pymongo.client_options import _parse_ssl_options
 from pymongo.common import (
     SRV_SERVICE_NAME,
     get_validated_options, INTERNAL_URI_OPTION_NAME_MAP,
@@ -569,6 +570,39 @@ def parse_uri(uri, default_port=DEFAULT_PORT, validate=True, warn=False,
     }
 
 
+def _parse_kms_tls_options(kms_tls_options):
+    """Parse KMS TLS connection options."""
+    if not kms_tls_options:
+        return {}
+    if not isinstance(kms_tls_options, dict):
+        raise TypeError('kms_tls_options must be a dict')
+    contexts = {}
+    for provider, opts in kms_tls_options.items():
+        if not isinstance(opts, dict):
+            raise TypeError(f'kms_tls_options["{provider}"] must be a dict')
+        opts.setdefault('tls', True)
+        opts = _CaseInsensitiveDictionary(opts)
+        opts = _handle_security_options(opts)
+        opts = _normalize_options(opts)
+        opts = validate_options(opts)
+        ssl_context, allow_invalid_hostnames = _parse_ssl_options(opts)
+        if ssl_context is None:
+            raise ConfigurationError('TLS is required for KMS providers')
+        if allow_invalid_hostnames:
+            raise ConfigurationError('Insecure TLS options prohibited')
+
+        for n in ['tlsInsecure',
+                  'tlsAllowInvalidCertificates',
+                  'tlsAllowInvalidHostnames',
+                  'tlsDisableOCSPEndpointCheck',
+                  'tlsDisableCertificateRevocationCheck']:
+            if n in opts:
+                raise ConfigurationError(
+                    f'Insecure TLS options prohibited: {n}')
+            contexts[provider] = ssl_context
+    return contexts
+
+
 if __name__ == '__main__':
     import pprint
     import sys
diff --git a/test/test_encryption.py b/test/test_encryption.py
@@ -17,6 +17,7 @@
 import base64
 import copy
 import os
+import ssl
 import traceback
 import socket
 import sys
@@ -50,9 +51,8 @@
 from pymongo.mongo_client import MongoClient
 from pymongo.operations import InsertOne
 from pymongo.write_concern import WriteConcern
-from test.test_ssl import CA_PEM
 
-from test import (unittest,
+from test import (unittest, CA_PEM, CLIENT_PEM,
                   client_context,
                   IntegrationTest,
                   PyMongoTestCase)
@@ -92,6 +92,7 @@ def test_init(self):
         self.assertEqual(opts._mongocryptd_spawn_path, 'mongocryptd')
         self.assertEqual(
             opts._mongocryptd_spawn_args, ['--idleShutdownTimeoutSecs=60'])
+        self.assertEqual(opts._kms_ssl_contexts, {})
 
     @unittest.skipUnless(_HAVE_PYMONGOCRYPT, 'pymongocrypt is not installed')
     def test_init_spawn_args(self):
@@ -116,6 +117,46 @@ def test_init_spawn_args(self):
             opts._mongocryptd_spawn_args,
             ['--quiet', '--port=27020', '--idleShutdownTimeoutSecs=60'])
 
+    @unittest.skipUnless(_HAVE_PYMONGOCRYPT, 'pymongocrypt is not installed')
+    def test_init_kms_tls_options(self):
+        # Error cases:
+        with self.assertRaisesRegex(
+                TypeError, r'kms_tls_options\["kmip"\] must be a dict'):
+            AutoEncryptionOpts({}, 'k.d', kms_tls_options={'kmip': 1})
+        for tls_opts in [
+                {'kmip': {'tls': True, 'tlsInsecure': True}},
+                {'kmip': {'tls': True, 'tlsAllowInvalidCertificates': True}},
+                {'kmip': {'tls': True, 'tlsAllowInvalidHostnames': True}},
+                {'kmip': {'tls': True, 'tlsDisableOCSPEndpointCheck': True}}]:
+            with self.assertRaisesRegex(
+                    ConfigurationError, 'Insecure TLS options prohibited'):
+                opts = AutoEncryptionOpts({}, 'k.d', kms_tls_options=tls_opts)
+        with self.assertRaises(FileNotFoundError):
+            AutoEncryptionOpts({}, 'k.d', kms_tls_options={
+                'kmip': {'tlsCAFile': 'does-not-exist'}})
+        # Success cases:
+        for tls_opts in [None, {}]:
+            opts = AutoEncryptionOpts({}, 'k.d', kms_tls_options=tls_opts)
+            self.assertEqual(opts._kms_ssl_contexts, {})
+        opts = AutoEncryptionOpts(
+            {}, 'k.d', kms_tls_options={'kmip': {'tls': True}, 'aws': {}})
+        ctx = opts._kms_ssl_contexts['kmip']
+        # On < 3.7 we check hostnames manually.
+        if sys.version_info[:2] >= (3, 7):
+            self.assertEqual(ctx.check_hostname, True)
+        self.assertEqual(ctx.verify_mode, ssl.CERT_REQUIRED)
+        ctx = opts._kms_ssl_contexts['aws']
+        if sys.version_info[:2] >= (3, 7):
+            self.assertEqual(ctx.check_hostname, True)
+        self.assertEqual(ctx.verify_mode, ssl.CERT_REQUIRED)
+        opts = AutoEncryptionOpts(
+            {}, 'k.d', kms_tls_options={'kmip': {
+                'tlsCAFile': CA_PEM, 'tlsCertificateKeyFile': CLIENT_PEM}})
+        ctx = opts._kms_ssl_contexts['kmip']
+        if sys.version_info[:2] >= (3, 7):
+            self.assertEqual(ctx.check_hostname, True)
+        self.assertEqual(ctx.verify_mode, ssl.CERT_REQUIRED)
+
 
 class TestClientOptions(PyMongoTestCase):
     def test_default(self):
