NO-SNOW: CRL changes in aio

Author: sfc-gh-turbaszek

src/snowflake/connector/aio/_connection.py

@@ -47,6 +47,7 @@
     PARAMETER_TIMEZONE,
     QueryStatus,
 )
+from ..crl import CRLConfig
 from ..description import PLATFORM, PYTHON_VERSION, SNOWFLAKE_CONNECTOR_VERSION
 from ..errorcode import (
     ER_CONNECTION_IS_CLOSED,
@@ -622,6 +623,7 @@ def _init_connection_parameters(
 
         # Placeholder attributes; will be initialized in connect()
         self._http_config: AioHttpConfig | None = None
+        self._crl_config: CRLConfig | None = None
         self._session_manager: SessionManager | None = None
         self._rest = None
         for name, (value, _) in DEFAULT_CONFIGURATION.items():

src/snowflake/connector/aio/_session_manager.py

@@ -3,15 +3,23 @@
 import sys
 from typing import TYPE_CHECKING
 
+import certifi
 from aiohttp import ClientRequest, ClientTimeout
 from aiohttp.client import _RequestOptions
 from aiohttp.client_proto import ResponseHandler
 from aiohttp.connector import Connection
 from aiohttp.typedefs import StrOrURL
 
 from .. import OperationalError
+from ..crl import CertRevocationCheckMode, CRLValidator
 from ..errorcode import ER_OCSP_RESPONSE_CERT_STATUS_REVOKED
-from ..ssl_wrap_socket import FEATURE_OCSP_RESPONSE_CACHE_FILE_NAME
+from ..ssl_wrap_socket import (
+    FEATURE_CRL_CONFIG,
+    FEATURE_OCSP_RESPONSE_CACHE_FILE_NAME,
+    get_current_session_manager,
+    load_trusted_certificates,
+    resolve_cafile,
+)
 from ._ocsp_asn1crypto import SnowflakeOCSPAsn1Crypto
 
 if TYPE_CHECKING:
@@ -71,6 +79,20 @@ async def connect(
     ) -> Connection:
         connection = await super().connect(req, traces, timeout)
         protocol = connection.protocol
+
+        logger.debug(
+            "CRL Check Mode: %s",
+            FEATURE_CRL_CONFIG.cert_revocation_check_mode.name,
+        )
+        if (
+            FEATURE_CRL_CONFIG.cert_revocation_check_mode
+            != CertRevocationCheckMode.ENABLED
+        ):
+            self.validate_crl(protocol, req)
+            logger.debug(
+                "The certificate revocation check was successful. No additional checks will be performed."
+            )
+
         if (
             req.is_ssl()
             and protocol is not None
@@ -90,6 +112,24 @@ async def connect(
                 protocol._snowflake_ocsp_validated = True
         return connection
 
+    def validate_crl(self, protocol: ResponseHandler, req: ClientRequest):
+        # Resolve CA file path from environment variables or use certifi default
+        cafile_for_ctx = resolve_cafile({"ca_certs": certifi.where()})
+        crl_validator = CRLValidator.from_config(
+            FEATURE_CRL_CONFIG,
+            get_current_session_manager(),
+            trusted_certificates=load_trusted_certificates(cafile_for_ctx),
+        )
+        sll_object = protocol.transport.get_extra_info("ssl_object")
+        if not crl_validator.validate_connection(sll_object):
+            raise OperationalError(
+                msg=(
+                    "The certificate is revoked or "
+                    "could not be validated via CRL: hostname={}".format(req.url.host)
+                ),
+                errno=ER_OCSP_RESPONSE_CERT_STATUS_REVOKED,
+            )
+
     async def validate_ocsp(
         self,
         hostname: str,

src/snowflake/connector/crl.py

@@ -1,12 +1,14 @@
 #!/usr/bin/env python
 from __future__ import annotations
 
+import ssl
 from collections import defaultdict
 from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
 from enum import Enum, unique
 from logging import getLogger
 from pathlib import Path
+from ssl import SSLObject
 from typing import Any
 
 from cryptography import x509
@@ -750,7 +752,7 @@ def validate_connection(self, connection: SSLConnection) -> bool:
         """
         try:
             # Get the peer certificate (the start certificate)
-            peer_cert = connection.get_peer_certificate(as_cryptography=True)
+            peer_cert = self._get_peer_certificate(connection)
             if peer_cert is None:
                 logger.warning("No peer certificate found in connection")
                 return (
@@ -765,6 +767,14 @@ def validate_connection(self, connection: SSLConnection) -> bool:
             logger.warning("Failed to validate connection: %s", e)
             return self._cert_revocation_check_mode == CertRevocationCheckMode.ADVISORY
 
+    @staticmethod
+    def _get_peer_certificate(connection):
+        if isinstance(connection, SSLObject):
+            return x509.load_der_x509_certificate(
+                connection.getpeercert(binary_form=True), default_backend()
+            )
+        return connection.get_peer_certificate(as_cryptography=True)
+
     def _extract_certificate_chain_from_connection(
         self, connection
     ) -> list[x509.Certificate] | None:
@@ -777,8 +787,17 @@ def _extract_certificate_chain_from_connection(
             Certificate chain as a list of x509.Certificate objects, or None on error
         """
         try:
-            # Convert OpenSSL certificates to cryptography x509 certificates
-            cert_chain = connection.get_peer_cert_chain(as_cryptography=True)
+            if isinstance(connection, SSLObject):
+                cert_chain = []
+                for cert in connection._sslobj.get_unverified_chain():
+                    cert_bytes = ssl.PEM_cert_to_DER_cert(cert.public_bytes())
+                    cert_chain.append(
+                        x509.load_der_x509_certificate(cert_bytes, default_backend())
+                    )
+            else:
+                # Convert OpenSSL certificates to cryptography x509 certificates
+                cert_chain = connection.get_peer_cert_chain(as_cryptography=True)
+
             if not cert_chain:
                 logger.debug("No certificate chain found in connection")
                 return None

src/snowflake/connector/ssl_wrap_socket.py

@@ -47,7 +47,7 @@
 
 
 # Helper utilities (private)
-def _resolve_cafile(kwargs: dict[str, Any]) -> str | None:
+def resolve_cafile(kwargs: dict[str, Any]) -> str | None:
     """Resolve CA bundle path from kwargs or standard environment variables.
 
     Precedence:
@@ -150,7 +150,7 @@ def inject_into_urllib3() -> None:
     connection_.ssl_wrap_socket = ssl_wrap_socket_with_cert_revocation_checks
 
 
-def _load_trusted_certificates(cafile: str | None) -> list[x509.Certificate]:
+def load_trusted_certificates(cafile: str | None) -> list[x509.Certificate]:
     # Use default SSL context to load the CA file and get the certificates
     ctx = ssl.create_default_context()
     ctx.load_verify_locations(cafile=cafile)
@@ -177,7 +177,7 @@ def ssl_wrap_socket_with_cert_revocation_checks(
 
     # Ensure PyOpenSSL context with partial-chain is used if none or wrong type provided
     provided_ctx = params.get("ssl_context")
-    cafile_for_ctx = _resolve_cafile(params)
+    cafile_for_ctx = resolve_cafile(params)
     if not isinstance(provided_ctx, PyOpenSSLContext):
         params["ssl_context"] = _build_context_with_partial_chain(cafile_for_ctx)
     else:
@@ -197,7 +197,7 @@ def ssl_wrap_socket_with_cert_revocation_checks(
         crl_validator = CRLValidator.from_config(
             FEATURE_CRL_CONFIG,
             get_current_session_manager(),
-            trusted_certificates=_load_trusted_certificates(cafile_for_ctx),
+            trusted_certificates=load_trusted_certificates(cafile_for_ctx),
         )
         if not crl_validator.validate_connection(ret.connection):
             raise OperationalError(