1 files changed, 47 insertions, 1 deletions
diff --git a/redis/connection.py b/redis/connection.py
index 4178f67..5fdac54 100755
--- a/redis/connection.py
+++ b/redis/connection.py
@@ -908,6 +908,9 @@ class SSLConnection(Connection):
         ssl_ca_path=None,
         ssl_password=None,
         ssl_validate_ocsp=False,
+        ssl_validate_ocsp_stapled=False,
+        ssl_ocsp_context=None,
+        ssl_ocsp_expected_cert=None,
         **kwargs,
     ):
         """Constructor
@@ -921,6 +924,11 @@ class SSLConnection(Connection):
             ssl_ca_path: The path to a directory containing several CA certificates in PEM format. Defaults to None.
             ssl_password: Password for unlocking an encrypted private key. Defaults to None.
 
+            ssl_validate_ocsp: If set, perform a full ocsp validation (i.e not a stapled verification)
+            ssl_validate_ocsp_stapled: If set, perform a validation on a stapled ocsp response
+            ssl_ocsp_context: A fully initialized OpenSSL.SSL.Context object to be used in verifying the ssl_ocsp_expected_cert
+            ssl_ocsp_expected_cert: A PEM armoured string containing the expected certificate to be returned from the ocsp verification service.
+
         Raises:
             RedisError
         """  # noqa
@@ -950,6 +958,9 @@ class SSLConnection(Connection):
         self.check_hostname = ssl_check_hostname
         self.certificate_password = ssl_password
         self.ssl_validate_ocsp = ssl_validate_ocsp
+        self.ssl_validate_ocsp_stapled = ssl_validate_ocsp_stapled
+        self.ssl_ocsp_context = ssl_ocsp_context
+        self.ssl_ocsp_expected_cert = ssl_ocsp_expected_cert
 
     def _connect(self):
         "Wrap the socket with SSL support"
@@ -968,7 +979,42 @@ class SSLConnection(Connection):
         sslsock = context.wrap_socket(sock, server_hostname=self.host)
         if self.ssl_validate_ocsp is True and CRYPTOGRAPHY_AVAILABLE is False:
             raise RedisError("cryptography is not installed.")
-        elif self.ssl_validate_ocsp is True and CRYPTOGRAPHY_AVAILABLE:
+
+        if self.ssl_validate_ocsp_stapled and self.ssl_validate_ocsp:
+            raise RedisError(
+                "Either an OCSP staple or pure OCSP connection must be validated "
+                "- not both."
+            )
+
+        # validation for the stapled case
+        if self.ssl_validate_ocsp_stapled:
+            import OpenSSL
+
+            from .ocsp import ocsp_staple_verifier
+
+            # if a context is provided use it - otherwise, a basic context
+            if self.ssl_ocsp_context is None:
+                staple_ctx = OpenSSL.SSL.Context(OpenSSL.SSL.SSLv23_METHOD)
+                staple_ctx.use_certificate_file(self.certfile)
+                staple_ctx.use_privatekey_file(self.keyfile)
+            else:
+                staple_ctx = self.ssl_ocsp_context
+
+            staple_ctx.set_ocsp_client_callback(
+                ocsp_staple_verifier,
+                self.ssl_ocsp_expected_cert,
+            )
+
+            #  need another socket
+            con = OpenSSL.SSL.Connection(staple_ctx, socket.socket())
+            con.request_ocsp()
+            con.connect((self.host, self.port))
+            con.do_handshake()
+            con.shutdown()
+            return sslsock
+
+        # pure ocsp validation
+        if self.ssl_validate_ocsp is True and CRYPTOGRAPHY_AVAILABLE:
             from .ocsp import OCSPVerifier
 
             o = OCSPVerifier(sslsock, self.host, self.port, self.ca_certs)