Fix error in ssl_wrap_socket_with_ocsp

Author: sfc-gh-pczajka

src/snowflake/connector/ssl_wrap_socket.py

@@ -104,19 +104,45 @@ def ssl_wrap_socket_with_ocsp(*args: Any, **kwargs: Any) -> WrappedSocket:
     ssl_context_index = get_args(ssl_.ssl_wrap_socket).args.index("ssl_context")
     context_in_args = len(args) > ssl_context_index
     ssl_context = (
-        args[hostname_index] if context_in_args else kwargs.get("ssl_context", None)
+        args[ssl_context_index] if context_in_args else kwargs.get("ssl_context", None)
     )
-    if not isinstance(ssl_context, PyOpenSSLContext):
-        # Create new default context
-        if context_in_args:
-            new_args = list(args)
-            new_args[ssl_context_index] = None
-            args = tuple(new_args)
-        else:
-            del kwargs["ssl_context"]
+
     # Fix ca certs location
     ca_certs_index = get_args(ssl_.ssl_wrap_socket).args.index("ca_certs")
     ca_certs_in_args = len(args) > ca_certs_index
+
+    if not isinstance(ssl_context, PyOpenSSLContext):
+        if FEATURE_OCSP_MODE != OCSPMode.DISABLE_OCSP_CHECKS:
+            # Create PyOpenSSL context for OCSP validation
+            ssl_context = PyOpenSSLContext(ssl_.ssl.PROTOCOL_TLS_CLIENT)
+            ssl_context.check_hostname = False
+            ssl_context.verify_mode = ssl_.ssl.CERT_REQUIRED
+
+            # Load CA certificates
+            ca_certs = (
+                kwargs.get("ca_certs")
+                or (args[ca_certs_index] if ca_certs_in_args else None)
+                or certifi.where()
+            )
+            ssl_context.load_verify_locations(ca_certs)
+
+            # Set the PyOpenSSL context in arguments
+            if context_in_args:
+                new_args = list(args)
+                new_args[ssl_context_index] = ssl_context
+                args = tuple(new_args)
+            else:
+                kwargs["ssl_context"] = ssl_context
+        else:
+            # Create new default context
+            if context_in_args:
+                new_args = list(args)
+                new_args[ssl_context_index] = None
+                args = tuple(new_args)
+            else:
+                if "ssl_context" in kwargs:
+                    del kwargs["ssl_context"]
+
     if not ca_certs_in_args and not kwargs.get("ca_certs"):
         kwargs["ca_certs"] = certifi.where()
 

src/snowflake/connector/vendored/urllib3/contrib/emscripten/__init__.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-import urllib3.connection
+import snowflake.connector.vendored.urllib3.connection
 
 from ...connectionpool import HTTPConnectionPool, HTTPSConnectionPool
 from .connection import EmscriptenHTTPConnection, EmscriptenHTTPSConnection
@@ -12,5 +12,5 @@ def inject_into_urllib3() -> None:
     # if it isn't ignored
     HTTPConnectionPool.ConnectionCls = EmscriptenHTTPConnection
     HTTPSConnectionPool.ConnectionCls = EmscriptenHTTPSConnection
-    urllib3.connection.HTTPConnection = EmscriptenHTTPConnection  # type: ignore[misc,assignment]
-    urllib3.connection.HTTPSConnection = EmscriptenHTTPSConnection  # type: ignore[misc,assignment]
+    snowflake.connector.vendored.urllib3.connection.HTTPConnection = EmscriptenHTTPConnection  # type: ignore[misc,assignment]
+    snowflake.connector.vendored.urllib3.connection.HTTPSConnection = EmscriptenHTTPSConnection  # type: ignore[misc,assignment]