add proxy support to OCSP

Author: smtakeda

chunk_downloader.py

@@ -12,6 +12,7 @@
 from .errorcode import (ER_NO_ADDITIONAL_CHUNK, ER_CHUNK_DOWNLOAD_FAILED)
 from .errors import (Error, OperationalError)
 from .network import (SnowflakeRestful, NO_TOKEN, MAX_CONNECTION_POOL)
+from .ssl_wrap_socket import (set_proxies)
 
 DEFAULT_REQUEST_TIMEOUT = 3600
 DEFAULT_CLIENT_RESULT_PREFETCH_SLOTS = 2
@@ -262,7 +263,7 @@ def _get_request(
         GET request for Large Result set chunkloader
         """
         # sharing the proxy and certificate
-        proxies = SnowflakeRestful.set_proxies(
+        proxies = set_proxies(
             self._connection.rest._proxy_host,
             self._connection.rest._proxy_port,
             self._connection.rest._proxy_user,

network.py

@@ -42,6 +42,7 @@
 from .sqlstate import (SQLSTATE_CONNECTION_NOT_EXISTS,
                        SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
                        SQLSTATE_CONNECTION_REJECTED)
+from .ssl_wrap_socket import (set_proxies)
 from .util_text import split_rows_from_stream
 from .version import VERSION
 
@@ -151,10 +152,18 @@ def __init__(self, host=u'127.0.0.1', port=8080,
         self._max_connection_pool = max_connection_pool
         self._connection = connection
         self.logger = getLogger(__name__)
+
+        # insecure mode (disabled by default)
         ssl_wrap_socket.FEATURE_INSECURE_MODE = \
             self._connection and self._connection._insecure_mode
+        # cache file name (enabled by default)
         ssl_wrap_socket.FEATURE_OCSP_RESPONSE_CACHE_FILE_NAME = \
             self._connection and self._connection._ocsp_response_cache_filename
+        #
+        ssl_wrap_socket.PROXY_HOST = self._proxy_host
+        ssl_wrap_socket.PROXY_PORT = self._proxy_port
+        ssl_wrap_socket.PROXY_USER = self._proxy_user
+        ssl_wrap_socket.PROXY_PASSWORD = self._proxy_password
 
         # This is to address the issue where requests hangs
         _ = 'dummy'.encode('idna').decode('utf-8')
@@ -168,35 +177,6 @@ def token(self):
     def master_token(self):
         return self._master_token if hasattr(self, u'_master_token') else None
 
-    @staticmethod
-    def set_proxies(proxy_host,
-                    proxy_port,
-                    proxy_user=None,
-                    proxy_password=None):
-        proxies = None
-        if proxy_host and proxy_port:
-            if proxy_user or proxy_password:
-                proxy_auth = u'{proxy_user}:{proxy_password}@'.format(
-                    proxy_user=proxy_user if proxy_user is not None else '',
-                    proxy_password=proxy_password if proxy_password is not
-                                                     None else ''
-                )
-            else:
-                proxy_auth = u''
-            proxies = {
-                u'http': u'http://{proxy_auth}{proxy_host}:{proxy_port}'.format(
-                    proxy_host=proxy_host,
-                    proxy_port=TO_UNICODE(proxy_port),
-                    proxy_auth=proxy_auth,
-                ),
-                u'https': u'http://{proxy_auth}{proxy_host}:{proxy_port}'.format(
-                    proxy_host=proxy_host,
-                    proxy_port=TO_UNICODE(proxy_port),
-                    proxy_auth=proxy_auth,
-                ),
-            }
-        return proxies
-
     def close(self):
         if hasattr(self, u'_token'):
             del self._token
@@ -509,7 +489,7 @@ def _get_request(self, url, headers, token=None, timeout=None):
             port=self._port,
             url=url,
         )
-        proxies = SnowflakeRestful.set_proxies(
+        proxies = set_proxies(
             self._proxy_host, self._proxy_port, self._proxy_user,
             self._proxy_password
         )
@@ -544,7 +524,7 @@ def _post_request(self, url, headers, body, token=None,
             port=self._port,
             url=url,
         )
-        proxies = SnowflakeRestful.set_proxies(
+        proxies = set_proxies(
             self._proxy_host, self._proxy_port, self._proxy_user,
             self._proxy_password)
 
@@ -883,7 +863,7 @@ def authenticate_by_saml(self, authenticator, account, user, password):
         token_url = data[u'tokenUrl']
         sso_url = data[u'ssoUrl']
 
-        proxies = SnowflakeRestful.set_proxies(
+        proxies = set_proxies(
             self._proxy_host, self._proxy_port, self._proxy_user,
             self._proxy_password)
         self.logger.debug(u'token_url=%s, proxies=%s', token_url, proxies)

ocsp_pyopenssl.py

@@ -513,7 +513,7 @@ def process_ocsp_response(response, ocsp_issuer):
     return single_response_map
 
 
-def execute_ocsp_request(ocsp_uri, cert_id, do_retry=True):
+def execute_ocsp_request(ocsp_uri, cert_id, proxies=None, do_retry=True):
     """
     Executes OCSP request for the given cert id
     """
@@ -558,6 +558,7 @@ def execute_ocsp_request(ocsp_uri, cert_id, do_retry=True):
                 'Host': parsed_url.hostname.encode(
                     'utf-8'),
             },
+            proxies=proxies,
             data=data)
         if response.status_code == OK:
             logger.debug("OCSP response was successfully returned")
@@ -1014,13 +1015,15 @@ class SnowflakeOCSP(object):
     OCSP validator using PyOpenSSL.
     """
 
-    def __init__(self, must_use_cache=False, ocsp_response_cache_url=None):
+    def __init__(self, must_use_cache=False,
+                 proxies=None, ocsp_response_cache_url=None):
         """
         :param must_use_cache: Test purpose. must use cache or raises an error
         :param ocsp_response_cache_url: the location of cache file
         """
         self.logger = getLogger(__name__)
         self._must_use_cache = must_use_cache
+        self._proxies = proxies
         if ocsp_response_cache_url is None and CACHE_DIR is not None:
             self._ocsp_response_cache_url = 'file://' + path.join(
                 CACHE_DIR, 'ocsp_response_cache')
@@ -1148,8 +1151,10 @@ def validate_by_direct_connection(
                 if not cache_status:
                     # not cached or invalid
                     self.logger.info('getting OCSP response from remote')
-                    ocsp_response = execute_ocsp_request(ocsp_uri, cert_id,
-                                                         do_retry=do_retry)
+                    ocsp_response = execute_ocsp_request(
+                        ocsp_uri, cert_id,
+                        proxies=self._proxies,
+                        do_retry=do_retry)
                 else:
                     self.logger.info('using OCSP response cache')
                 single_response_map = process_ocsp_response(
@@ -1175,7 +1180,8 @@ def validate_by_direct_connection(
 
         return True, cert_id, ocsp_response
 
-    def generate_cert_id_response(self, hostname, connection, do_retry=True):
+    def generate_cert_id_response(
+            self, hostname, connection, proxies=None, do_retry=True):
         current_time = int(time.time())
         cert_data = _extract_certificate_chain(connection)
         results = {}
@@ -1188,7 +1194,8 @@ def generate_cert_id_response(self, hostname, connection, do_retry=True):
             if ocsp_uri:
                 ret, cert_id, ocsp_response = \
                     self.validate_by_direct_connection(
-                        ocsp_uri, ocsp_issuer, ocsp_subject, do_retry)
+                        ocsp_uri, ocsp_issuer, ocsp_subject,
+                        do_retry=do_retry)
                 if ret and cert_id and ocsp_response:
                     cert_id_der = der_encoder.encode(cert_id)
                     results[cert_id_der] = (
@@ -1242,7 +1249,7 @@ def _openssl_connect(hostname, port=443):
 
     ocsp = SnowflakeOCSP()
     connection = _openssl_connect(url, port)
-    results = ocsp.generate_cert_id_response(url, connection)
+    results = ocsp.generate_cert_id_response(url, connection, proxies=None)
     current_Time = int(time.time())
     print("Target URL: https://{0}:{1}/".format(url, port))
     print("Current Time: {0}".format(

ssl_wrap_socket.py

@@ -18,6 +18,14 @@
 """
 FEATURE_OCSP_RESPONSE_CACHE_FILE_NAME = None
 
+"""
+Proxy, shared across all connections
+"""
+PROXY_HOST = None
+PROXY_PORT = None
+PROXY_USER = None
+PROXY_PASSWORD = None
+
 # imports
 import select
 import socket
@@ -36,6 +44,7 @@
 from cryptography.hazmat.backends.openssl import backend as openssl_backend
 from cryptography.hazmat.backends.openssl.x509 import _Certificate
 
+from .compat import (TO_UNICODE)
 from .errorcode import (ER_SERVER_CERTIFICATE_REVOKED)
 from .errors import (OperationalError)
 from .ocsp_pyopenssl import SnowflakeOCSP
@@ -414,6 +423,8 @@ def ssl_wrap_socket_with_ocsp(
                 FEATURE_OCSP_RESPONSE_CACHE_FILE_NAME)
     if not FEATURE_INSECURE_MODE:
         v = SnowflakeOCSP(
+            proxies=set_proxies(PROXY_HOST, PROXY_PORT, PROXY_USER,
+                                PROXY_PASSWORD),
             ocsp_response_cache_url=FEATURE_OCSP_RESPONSE_CACHE_FILE_NAME
         ).validate(server_hostname, ret.connection)
         if not v:
@@ -432,6 +443,38 @@ def ssl_wrap_socket_with_ocsp(
     return ret
 
 
+def set_proxies(proxy_host,
+                proxy_port,
+                proxy_user=None,
+                proxy_password=None):
+    """
+    Set proxy dict for requests
+    """
+    proxies = None
+    if proxy_host and proxy_port:
+        if proxy_user or proxy_password:
+            proxy_auth = u'{proxy_user}:{proxy_password}@'.format(
+                proxy_user=proxy_user if proxy_user is not None else '',
+                proxy_password=proxy_password if proxy_password is not
+                                                 None else ''
+            )
+        else:
+            proxy_auth = u''
+        proxies = {
+            u'http': u'http://{proxy_auth}{proxy_host}:{proxy_port}'.format(
+                proxy_host=proxy_host,
+                proxy_port=TO_UNICODE(proxy_port),
+                proxy_auth=proxy_auth,
+            ),
+            u'https': u'http://{proxy_auth}{proxy_host}:{proxy_port}'.format(
+                proxy_host=proxy_host,
+                proxy_port=TO_UNICODE(proxy_port),
+                proxy_auth=proxy_auth,
+            ),
+        }
+    return proxies
+
+
 def inject_into_urllib3():
     """
     Monkey-patch urllib3 with PyOpenSSL-backed SSL-support and OCSP