move kms_ssl_contexts

Author: sleepyStick

pymongo/asynchronous/encryption.py

@@ -87,7 +87,7 @@
 from pymongo.results import BulkWriteResult, DeleteResult
 from pymongo.ssl_support import BLOCKING_IO_ERRORS, get_ssl_context
 from pymongo.typings import _DocumentType, _DocumentTypeArg
-from pymongo.uri_parser_shared import parse_host
+from pymongo.uri_parser_shared import _parse_kms_tls_options, parse_host
 from pymongo.write_concern import WriteConcern
 
 if TYPE_CHECKING:
@@ -157,6 +157,7 @@ def __init__(
         self.mongocryptd_client = mongocryptd_client
         self.opts = opts
         self._spawned = False
+        self._kms_ssl_contexts = _parse_kms_tls_options(opts._kms_tls_options, _IS_SYNC)
 
     async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
         """Complete a KMS request.
@@ -165,11 +166,10 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
 
         :return: None
         """
-        self.opts._parse_kms_tls_options(_IS_SYNC)
         endpoint = kms_context.endpoint
         message = kms_context.message
         provider = kms_context.kms_provider
-        ctx = self.opts._kms_ssl_contexts.get(provider)
+        ctx = self._kms_ssl_contexts.get(provider)
         if ctx is None:
             # Enable strict certificate verification, OCSP, match hostname, and
             # SNI using the system default CA certificates.
@@ -677,7 +677,7 @@ def __init__(
             kms_tls_options=kms_tls_options,
             key_expiration_ms=key_expiration_ms,
         )
-        opts._parse_kms_tls_options(_IS_SYNC)
+        self._kms_ssl_contexts = _parse_kms_tls_options(opts._kms_tls_options, _IS_SYNC)
         self._io_callbacks: Optional[_EncryptionIO] = _EncryptionIO(
             None, key_vault_coll, None, opts
         )

pymongo/encryption_options.py

@@ -32,10 +32,8 @@
 from bson import int64
 from pymongo.common import validate_is_mapping
 from pymongo.errors import ConfigurationError
-from pymongo.uri_parser_shared import _parse_kms_tls_options
 
 if TYPE_CHECKING:
-    from pymongo.pyopenssl_context import SSLContext
     from pymongo.typings import _AgnosticMongoClient, _DocumentTypeArg
 
 
@@ -238,13 +236,9 @@ def __init__(
             self._mongocryptd_spawn_args.append("--idleShutdownTimeoutSecs=60")
         # Maps KMS provider name to a SSLContext.
         self._kms_tls_options = kms_tls_options
-        self._kms_ssl_contexts: dict[str, SSLContext] = {}
         self._bypass_query_analysis = bypass_query_analysis
         self._key_expiration_ms = key_expiration_ms
 
-    def _parse_kms_tls_options(self, is_sync: bool) -> None:
-        self._kms_ssl_contexts = _parse_kms_tls_options(self._kms_tls_options, is_sync)
-
 
 class RangeOpts:
     """Options to configure encrypted queries using the range algorithm."""

pymongo/synchronous/encryption.py

@@ -86,7 +86,7 @@
 from pymongo.synchronous.database import Database
 from pymongo.synchronous.mongo_client import MongoClient
 from pymongo.typings import _DocumentType, _DocumentTypeArg
-from pymongo.uri_parser_shared import parse_host
+from pymongo.uri_parser_shared import _parse_kms_tls_options, parse_host
 from pymongo.write_concern import WriteConcern
 
 if TYPE_CHECKING:
@@ -156,6 +156,7 @@ def __init__(
         self.mongocryptd_client = mongocryptd_client
         self.opts = opts
         self._spawned = False
+        self._kms_ssl_contexts = _parse_kms_tls_options(opts._kms_tls_options, _IS_SYNC)
 
     def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
         """Complete a KMS request.
@@ -164,11 +165,10 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
 
         :return: None
         """
-        self.opts._parse_kms_tls_options(_IS_SYNC)
         endpoint = kms_context.endpoint
         message = kms_context.message
         provider = kms_context.kms_provider
-        ctx = self.opts._kms_ssl_contexts.get(provider)
+        ctx = self._kms_ssl_contexts.get(provider)
         if ctx is None:
             # Enable strict certificate verification, OCSP, match hostname, and
             # SNI using the system default CA certificates.
@@ -670,7 +670,7 @@ def __init__(
             kms_tls_options=kms_tls_options,
             key_expiration_ms=key_expiration_ms,
         )
-        opts._parse_kms_tls_options(_IS_SYNC)
+        self._kms_ssl_contexts = _parse_kms_tls_options(opts._kms_tls_options, _IS_SYNC)
         self._io_callbacks: Optional[_EncryptionIO] = _EncryptionIO(
             None, key_vault_coll, None, opts
         )

test/asynchronous/test_encryption.py

@@ -41,6 +41,7 @@
 from pymongo.asynchronous.collection import AsyncCollection
 from pymongo.asynchronous.helpers import anext
 from pymongo.daemon import _spawn_daemon
+from pymongo.uri_parser import _parse_kms_tls_options
 
 try:
     from pymongo.pyopenssl_context import IS_PYOPENSSL
@@ -141,7 +142,7 @@ def test_init(self):
         self.assertEqual(opts._mongocryptd_bypass_spawn, False)
         self.assertEqual(opts._mongocryptd_spawn_path, "mongocryptd")
         self.assertEqual(opts._mongocryptd_spawn_args, ["--idleShutdownTimeoutSecs=60"])
-        self.assertEqual(opts._kms_ssl_contexts, {})
+        self.assertEqual(opts._kms_tls_options, {})
 
     @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed")
     def test_init_spawn_args(self):
@@ -189,22 +190,22 @@ def test_init_kms_tls_options(self):
         tls_opts: Any
         for tls_opts in [None, {}]:
             opts = AutoEncryptionOpts({}, "k.d", kms_tls_options=tls_opts)
-            self.assertEqual(opts._kms_ssl_contexts, {})
+            self.assertEqual(opts._kms_tls_options, {})
         opts = AutoEncryptionOpts({}, "k.d", kms_tls_options={"kmip": {"tls": True}, "aws": {}})
-        opts._parse_kms_tls_options(_IS_SYNC)
-        ctx = opts._kms_ssl_contexts["kmip"]
+        _kms_ssl_contexts = _parse_kms_tls_options(opts._kms_tls_options, _IS_SYNC)
+        ctx = _kms_ssl_contexts["kmip"]
         self.assertEqual(ctx.check_hostname, True)
         self.assertEqual(ctx.verify_mode, ssl.CERT_REQUIRED)
-        ctx = opts._kms_ssl_contexts["aws"]
+        ctx = _kms_ssl_contexts["aws"]
         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}},
         )
-        opts._parse_kms_tls_options(_IS_SYNC)
-        ctx = opts._kms_ssl_contexts["kmip"]
+        _kms_ssl_contexts = _parse_kms_tls_options(opts._kms_tls_options, _IS_SYNC)
+        ctx = _kms_ssl_contexts["kmip"]
         self.assertEqual(ctx.check_hostname, True)
         self.assertEqual(ctx.verify_mode, ssl.CERT_REQUIRED)
 
@@ -2233,7 +2234,7 @@ async def test_05_tlsDisableOCSPEndpointCheck_is_permitted(self):
         encryption = self.create_client_encryption(
             providers, "keyvault.datakeys", self.client, OPTS, kms_tls_options=options
         )
-        ctx = encryption._io_callbacks.opts._kms_ssl_contexts["aws"]
+        ctx = encryption._io_callbacks._kms_ssl_contexts["aws"]
         if not hasattr(ctx, "check_ocsp_endpoint"):
             raise self.skipTest("OCSP not enabled")
         self.assertFalse(ctx.check_ocsp_endpoint)

test/test_encryption.py

@@ -41,6 +41,7 @@
 from pymongo.daemon import _spawn_daemon
 from pymongo.synchronous.collection import Collection
 from pymongo.synchronous.helpers import next
+from pymongo.uri_parser import _parse_kms_tls_options
 
 try:
     from pymongo.pyopenssl_context import IS_PYOPENSSL
@@ -141,7 +142,7 @@ def test_init(self):
         self.assertEqual(opts._mongocryptd_bypass_spawn, False)
         self.assertEqual(opts._mongocryptd_spawn_path, "mongocryptd")
         self.assertEqual(opts._mongocryptd_spawn_args, ["--idleShutdownTimeoutSecs=60"])
-        self.assertEqual(opts._kms_ssl_contexts, {})
+        self.assertEqual(opts._kms_tls_options, {})
 
     @unittest.skipUnless(_HAVE_PYMONGOCRYPT, "pymongocrypt is not installed")
     def test_init_spawn_args(self):
@@ -189,22 +190,22 @@ def test_init_kms_tls_options(self):
         tls_opts: Any
         for tls_opts in [None, {}]:
             opts = AutoEncryptionOpts({}, "k.d", kms_tls_options=tls_opts)
-            self.assertEqual(opts._kms_ssl_contexts, {})
+            self.assertEqual(opts._kms_tls_options, {})
         opts = AutoEncryptionOpts({}, "k.d", kms_tls_options={"kmip": {"tls": True}, "aws": {}})
-        opts._parse_kms_tls_options(_IS_SYNC)
-        ctx = opts._kms_ssl_contexts["kmip"]
+        _kms_ssl_contexts = _parse_kms_tls_options(opts._kms_tls_options, _IS_SYNC)
+        ctx = _kms_ssl_contexts["kmip"]
         self.assertEqual(ctx.check_hostname, True)
         self.assertEqual(ctx.verify_mode, ssl.CERT_REQUIRED)
-        ctx = opts._kms_ssl_contexts["aws"]
+        ctx = _kms_ssl_contexts["aws"]
         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}},
         )
-        opts._parse_kms_tls_options(_IS_SYNC)
-        ctx = opts._kms_ssl_contexts["kmip"]
+        _kms_ssl_contexts = _parse_kms_tls_options(opts._kms_tls_options, _IS_SYNC)
+        ctx = _kms_ssl_contexts["kmip"]
         self.assertEqual(ctx.check_hostname, True)
         self.assertEqual(ctx.verify_mode, ssl.CERT_REQUIRED)
 
@@ -2225,7 +2226,7 @@ def test_05_tlsDisableOCSPEndpointCheck_is_permitted(self):
         encryption = self.create_client_encryption(
             providers, "keyvault.datakeys", self.client, OPTS, kms_tls_options=options
         )
-        ctx = encryption._io_callbacks.opts._kms_ssl_contexts["aws"]
+        ctx = encryption._io_callbacks._kms_ssl_contexts["aws"]
         if not hasattr(ctx, "check_ocsp_endpoint"):
             raise self.skipTest("OCSP not enabled")
         self.assertFalse(ctx.check_ocsp_endpoint)