snowflake-connector-python 4.4.0__cp314-cp314-win_amd64.whl

snowflake/connector/ocsp_asn1crypto.py

@@ -0,0 +1,447 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import typing
+from base64 import b64decode, b64encode
+from collections import OrderedDict
+from datetime import datetime, timezone
+from logging import getLogger
+from os import getenv
+
+from asn1crypto.algos import DigestAlgorithm
+from asn1crypto.core import Integer, OctetString
+from asn1crypto.ocsp import (
+    CertId,
+    OCSPRequest,
+    OCSPResponse,
+    Request,
+    Requests,
+    SingleResponse,
+    TBSRequest,
+    Version,
+)
+from asn1crypto.x509 import Certificate
+from cryptography.exceptions import InvalidSignature
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import padding, utils
+from cryptography.hazmat.primitives.asymmetric.dsa import DSAPublicKey
+from cryptography.hazmat.primitives.asymmetric.ec import ECDSA, EllipticCurvePublicKey
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey
+from OpenSSL.SSL import Connection
+
+from snowflake.connector.errorcode import (
+    ER_OCSP_RESPONSE_ATTACHED_CERT_EXPIRED,
+    ER_OCSP_RESPONSE_ATTACHED_CERT_INVALID,
+    ER_OCSP_RESPONSE_CERT_STATUS_INVALID,
+    ER_OCSP_RESPONSE_INVALID_SIGNATURE,
+    ER_OCSP_RESPONSE_LOAD_FAILURE,
+    ER_OCSP_RESPONSE_STATUS_UNSUCCESSFUL,
+)
+from snowflake.connector.errors import RevocationCheckError
+from snowflake.connector.ocsp_snowflake import SnowflakeOCSP, generate_cache_key
+
+logger = getLogger(__name__)
+
+
+class SnowflakeOCSPAsn1Crypto(SnowflakeOCSP):
+    """OCSP checks by asn1crypto."""
+
+    # map signature algorithm name to digest class
+    SIGNATURE_ALGORITHM_TO_DIGEST_CLASS = {
+        "sha256": hashes.SHA256,
+        "sha384": hashes.SHA384,
+        "sha512": hashes.SHA512,
+    }
+
+    def encode_cert_id_key(self, hkey):
+        issuer_name_hash, issuer_key_hash, serial_number = hkey
+        issuer_name_hash = OctetString.load(issuer_name_hash)
+        issuer_key_hash = OctetString.load(issuer_key_hash)
+        serial_number = Integer.load(serial_number)
+        cert_id = CertId(
+            {
+                "hash_algorithm": DigestAlgorithm(
+                    {"algorithm": "sha1", "parameters": None}
+                ),
+                "issuer_name_hash": issuer_name_hash,
+                "issuer_key_hash": issuer_key_hash,
+                "serial_number": serial_number,
+            }
+        )
+        return cert_id
+
+    def decode_cert_id_key(self, cert_id: CertId) -> tuple[bytes, bytes, bytes]:
+        return generate_cache_key(cert_id)
+
+    def decode_cert_id_base64(self, cert_id_base64):
+        return CertId.load(b64decode(cert_id_base64))
+
+    def encode_cert_id_base64(self, hkey):
+        return b64encode(self.encode_cert_id_key(hkey).dump()).decode("ascii")
+
+    def read_cert_bundle(self, ca_bundle_file, storage=None) -> None:
+        """Reads a certificate file including certificates in PEM format."""
+        if storage is None:
+            storage = SnowflakeOCSP.ROOT_CERTIFICATES_DICT
+        logger.debug("reading certificate bundle: %s", ca_bundle_file)
+        with open(ca_bundle_file, "rb") as all_certs:
+            # don't lock storage
+            from asn1crypto import pem
+
+            pem_certs = pem.unarmor(all_certs.read(), multiple=True)
+            for type_name, _, der_bytes in pem_certs:
+                if type_name == "CERTIFICATE":
+                    crt = Certificate.load(der_bytes)
+                    storage[crt.subject.sha256] = crt
+
+    def create_ocsp_request(
+        self,
+        issuer: Certificate,
+        subject: Certificate,
+    ) -> tuple[CertId, OCSPRequest]:
+        """Creates CertId and OCSPRequest."""
+        cert_id = CertId(
+            {
+                "hash_algorithm": DigestAlgorithm(
+                    {"algorithm": "sha1", "parameters": None}
+                ),
+                "issuer_name_hash": OctetString(subject.issuer.sha1),
+                "issuer_key_hash": OctetString(issuer.public_key.sha1),
+                "serial_number": subject.serial_number,
+            }
+        )
+        ocsp_request = OCSPRequest(
+            {
+                "tbs_request": TBSRequest(
+                    {
+                        "version": Version(0),
+                        "request_list": Requests(
+                            [
+                                Request(
+                                    {
+                                        "req_cert": cert_id,
+                                    }
+                                )
+                            ]
+                        ),
+                    }
+                ),
+            }
+        )
+        return cert_id, ocsp_request
+
+    def extract_ocsp_url(self, cert):
+        urls = cert.ocsp_urls
+        ocsp_url = urls[0] if urls else None
+        return ocsp_url
+
+    def decode_ocsp_request(self, ocsp_request):
+        return ocsp_request.dump()
+
+    def decode_ocsp_request_b64(self, ocsp_request):
+        data = self.decode_ocsp_request(ocsp_request)  # convert to DER
+        b64data = b64encode(data).decode("ascii")
+        return b64data
+
+    def extract_good_status(
+        self, single_response: SingleResponse
+    ) -> tuple[datetime, datetime]:
+        """Extracts GOOD status."""
+        this_update_native = single_response["this_update"].native
+        next_update_native = single_response["next_update"].native
+
+        return this_update_native, next_update_native
+
+    def extract_revoked_status(self, single_response):
+        """Extracts REVOKED status."""
+        revoked_info = single_response["cert_status"]
+        revocation_time = revoked_info.native["revocation_time"]
+        revocation_reason = revoked_info.native["revocation_reason"]
+        return revocation_time, revocation_reason
+
+    def check_cert_time_validity(
+        self, cur_time: datetime, ocsp_cert: Certificate
+    ) -> tuple[bool, str | None]:
+        val_start = ocsp_cert["tbs_certificate"]["validity"]["not_before"].native
+        val_end = ocsp_cert["tbs_certificate"]["validity"]["not_after"].native
+
+        if cur_time > val_end or cur_time < val_start:
+            debug_msg = (
+                "Certificate attached to OCSP response is invalid. OCSP response "
+                "current time - {} certificate not before time - {} certificate "
+                "not after time - {}. Consider running curl -o ocsp.der {}".format(
+                    cur_time,
+                    val_start,
+                    val_end,
+                    super().debug_ocsp_failure_url,
+                )
+            )
+
+            return False, debug_msg
+        else:
+            return True, None
+
+    """
+    is_valid_time - checks various components of the OCSP Response
+    for expiry.
+    :param cert_id - certificate id corresponding to OCSP Response
+    :param ocsp_response
+    :return True/False depending on time validity within the response
+    """
+
+    def is_valid_time(self, cert_id, ocsp_response) -> bool:
+        res = OCSPResponse.load(ocsp_response)
+
+        if res["response_status"].native != "successful":
+            raise RevocationCheckError(
+                msg="Invalid Status: {}".format(res["response_status"].native),
+                errno=ER_OCSP_RESPONSE_STATUS_UNSUCCESSFUL,
+            )
+
+        basic_ocsp_response = res.basic_ocsp_response
+        if basic_ocsp_response["certs"].native:
+            ocsp_cert = basic_ocsp_response["certs"][0]
+            logger.debug(
+                "Verifying the attached certificate is signed by "
+                "the issuer. Valid Not After: %s",
+                ocsp_cert["tbs_certificate"]["validity"]["not_after"].native,
+            )
+
+            cur_time = datetime.now(timezone.utc)
+
+            """
+            Note:
+            We purposefully do not verify certificate signature here.
+            The OCSP Response is extracted from the OCSP Response Cache
+            which is expected to have OCSP Responses with verified
+            attached signature. Moreover this OCSP Response is eventually
+            going to be processed by the driver before being consumed by
+            the driver.
+            This step ensures that the OCSP Response cache does not have
+            any invalid entries.
+            """
+            cert_valid, debug_msg = self.check_cert_time_validity(cur_time, ocsp_cert)
+            if not cert_valid:
+                logger.debug(debug_msg)
+                return False
+
+        tbs_response_data = basic_ocsp_response["tbs_response_data"]
+
+        single_response = tbs_response_data["responses"][0]
+        cert_status = single_response["cert_status"].name
+
+        try:
+            if cert_status == "good":
+                self._process_good_status(single_response, cert_id, ocsp_response)
+        except Exception as ex:
+            logger.debug("Failed to validate ocsp response %s", ex)
+            return False
+
+        return True
+
+    def process_ocsp_response(self, issuer, cert_id, ocsp_response):
+        try:
+            res = OCSPResponse.load(ocsp_response)
+            if self.test_mode is not None:
+                ocsp_load_failure = getenv("SF_TEST_OCSP_FORCE_BAD_OCSP_RESPONSE")
+                if ocsp_load_failure is not None:
+                    raise RevocationCheckError(
+                        msg="Force fail", errno=ER_OCSP_RESPONSE_LOAD_FAILURE
+                    )
+        except Exception:
+            raise RevocationCheckError(
+                msg="Invalid OCSP Response", errno=ER_OCSP_RESPONSE_LOAD_FAILURE
+            )
+
+        if res["response_status"].native != "successful":
+            raise RevocationCheckError(
+                msg="Invalid Status: {}".format(res["response_status"].native),
+                errno=ER_OCSP_RESPONSE_STATUS_UNSUCCESSFUL,
+            )
+
+        basic_ocsp_response = res.basic_ocsp_response
+        if basic_ocsp_response["certs"].native:
+            logger.debug("Certificate is attached in Basic OCSP Response")
+            ocsp_cert = basic_ocsp_response["certs"][0]
+            logger.debug(
+                "Verifying the attached certificate is signed by " "the issuer"
+            )
+            logger.debug(
+                "Valid Not After: %s",
+                ocsp_cert["tbs_certificate"]["validity"]["not_after"].native,
+            )
+
+            cur_time = datetime.now(timezone.utc)
+
+            try:
+                """
+                Signature verification should happen before any kind of
+                validation
+                """
+                self.verify_signature(
+                    ocsp_cert.hash_algo,
+                    ocsp_cert.signature,
+                    issuer,
+                    ocsp_cert["tbs_certificate"],
+                )
+            except RevocationCheckError as rce:
+                raise RevocationCheckError(
+                    msg=rce.msg, errno=ER_OCSP_RESPONSE_ATTACHED_CERT_INVALID
+                )
+            cert_valid, debug_msg = self.check_cert_time_validity(cur_time, ocsp_cert)
+
+            if not cert_valid:
+                raise RevocationCheckError(
+                    msg=debug_msg, errno=ER_OCSP_RESPONSE_ATTACHED_CERT_EXPIRED
+                )
+
+        else:
+            logger.debug(
+                "Certificate is NOT attached in Basic OCSP Response. "
+                "Using issuer's certificate"
+            )
+            ocsp_cert = issuer
+
+        tbs_response_data = basic_ocsp_response["tbs_response_data"]
+
+        logger.debug("Verifying the OCSP response is signed by the issuer.")
+        try:
+            self.verify_signature(
+                basic_ocsp_response["signature_algorithm"].hash_algo,
+                basic_ocsp_response["signature"].native,
+                ocsp_cert,
+                tbs_response_data,
+            )
+        except RevocationCheckError as rce:
+            raise RevocationCheckError(
+                msg=rce.msg, errno=ER_OCSP_RESPONSE_INVALID_SIGNATURE
+            )
+
+        single_response = tbs_response_data["responses"][0]
+        cert_status = single_response["cert_status"].name
+        if self.test_mode is not None:
+            test_cert_status = getenv("SF_TEST_OCSP_CERT_STATUS")
+            if test_cert_status == "revoked":
+                cert_status = "revoked"
+            elif test_cert_status == "unknown":
+                cert_status = "unknown"
+            elif test_cert_status == "good":
+                cert_status = "good"
+
+        try:
+            if cert_status == "good":
+                self._process_good_status(single_response, cert_id, ocsp_response)
+            elif cert_status == "revoked":
+                self._process_revoked_status(single_response, cert_id)
+            elif cert_status == "unknown":
+                self._process_unknown_status(cert_id)
+            else:
+                debug_msg = (
+                    "Unknown revocation status was returned."
+                    "OCSP response may be malformed: {}.".format(cert_status)
+                )
+                raise RevocationCheckError(
+                    msg=debug_msg, errno=ER_OCSP_RESPONSE_CERT_STATUS_INVALID
+                )
+        except RevocationCheckError as op_er:
+            debug_msg = "{} Consider running curl -o ocsp.der {}".format(
+                op_er.msg, self.debug_ocsp_failure_url
+            )
+            raise RevocationCheckError(msg=debug_msg, errno=op_er.errno)
+
+    def verify_signature(self, signature_algorithm, signature, cert, data):
+        backend = default_backend()
+        public_key = serialization.load_der_public_key(
+            cert.public_key.dump(), backend=default_backend()
+        )
+        if (
+            signature_algorithm
+            in SnowflakeOCSPAsn1Crypto.SIGNATURE_ALGORITHM_TO_DIGEST_CLASS
+        ):
+            chosen_hash = SnowflakeOCSPAsn1Crypto.SIGNATURE_ALGORITHM_TO_DIGEST_CLASS[
+                signature_algorithm
+            ]()
+        else:
+            # the last resort. should not happen.
+            chosen_hash = hashes.SHA1()
+        hasher = hashes.Hash(chosen_hash, backend)
+        hasher.update(data.dump())
+        digest = hasher.finalize()
+        additional_kwargs: dict[str, typing.Any] = dict()
+        if isinstance(public_key, RSAPublicKey):
+            additional_kwargs["padding"] = padding.PKCS1v15()
+            additional_kwargs["algorithm"] = utils.Prehashed(chosen_hash)
+        elif isinstance(public_key, DSAPublicKey):
+            additional_kwargs["algorithm"] = utils.Prehashed(chosen_hash)
+        elif isinstance(public_key, EllipticCurvePublicKey):
+            additional_kwargs["signature_algorithm"] = ECDSA(
+                utils.Prehashed(chosen_hash)
+            )
+        try:
+            public_key.verify(
+                signature,
+                digest,
+                **additional_kwargs,
+            )
+        except InvalidSignature:
+            raise RevocationCheckError(msg="Failed to verify the signature")
+
+    def extract_certificate_chain(
+        self, connection: Connection
+    ) -> list[tuple[Certificate, Certificate]]:
+        """Gets certificate chain and extract the key info from OpenSSL connection."""
+        from OpenSSL.crypto import FILETYPE_ASN1, dump_certificate
+
+        cert_map = OrderedDict()
+        cert_chain = connection.get_peer_cert_chain()
+        logger.debug("# of certificates: %s", len(cert_chain))
+        self._lazy_read_ca_bundle()
+        for cert_openssl in cert_chain:
+            cert_der = dump_certificate(FILETYPE_ASN1, cert_openssl)
+            cert = Certificate.load(cert_der)
+            logger.debug(
+                "subject: %s, issuer: %s", cert.subject.native, cert.issuer.native
+            )
+            cert_map[cert.subject.sha256] = cert
+            if cert.issuer.sha256 in SnowflakeOCSP.ROOT_CERTIFICATES_DICT:
+                logger.debug(
+                    "A trusted root certificate found: %s, stopping chain traversal here",
+                    cert.subject.native,
+                )
+                break
+
+        return self.create_pair_issuer_subject(cert_map)
+
+    def create_pair_issuer_subject(
+        self, cert_map: OrderedDict
+    ) -> list[tuple[Certificate, Certificate]]:
+        """Creates pairs of issuer and subject certificates."""
+        issuer_subject = []
+        for subject_der in cert_map:
+            subject = cert_map[subject_der]
+            if subject.ocsp_no_check_value or subject.ca and not subject.ocsp_urls:
+                # Root certificate will not be validated
+                # but it is used to validate the subject certificate
+                continue
+            issuer_hash = subject.issuer.sha256
+            if issuer_hash not in cert_map:
+                # IF NO ROOT certificate is attached in the certificate chain
+                # read it from the local disk
+                self._lazy_read_ca_bundle()
+                logger.debug("not found issuer_der: %s", subject.issuer.native)
+                if issuer_hash not in SnowflakeOCSP.ROOT_CERTIFICATES_DICT:
+                    raise RevocationCheckError(
+                        msg="CA certificate is NOT found in the root "
+                        "certificate list. Make sure you use the latest "
+                        "Python Connector package and the URL is valid."
+                    )
+                issuer = SnowflakeOCSP.ROOT_CERTIFICATES_DICT[issuer_hash]
+            else:
+                issuer = cert_map[issuer_hash]
+
+            issuer_subject.append((issuer, subject))
+        return issuer_subject
+
+    def subject_name(self, subject: Certificate) -> OrderedDict:
+        return subject.subject.native