Merge pull request #18 from Juniper/proxy_support

Adding proxy support for HBEZ
Juniper · Aug 4, 2021 · cf2430c · cf2430c
2 parents 99a9c54 + 514438f
commit cf2430c
Showing 1 changed file with 24 additions and 14 deletions.
diff --git a/lib/jnpr/healthbot/swagger/rest.py b/lib/jnpr/healthbot/swagger/rest.py
@@ -10,7 +10,6 @@
     Generated by: https://github.com/swagger-api/swagger-codegen.git
 """
 
-
 from __future__ import absolute_import
 
 import io
@@ -29,7 +28,6 @@
 except ImportError:
     raise ImportError('Swagger python client requires urllib3.')
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -84,16 +82,28 @@ def __init__(self, configuration, pools_size=4, maxsize=None):
 
         # https pool manager
         if configuration.proxy:
-            self.pool_manager = urllib3.ProxyManager(
-                num_pools=pools_size,
-                maxsize=maxsize,
-                cert_reqs=cert_reqs,
-                ca_certs=ca_certs,
-                cert_file=configuration.cert_file,
-                key_file=configuration.key_file,
-                proxy_url=configuration.proxy,
-                **addition_pool_args
-            )
+            if 'socks' in configuration.proxy:
+                self.pool_manager = urllib3.contrib.socks.SOCKSProxyManager(
+                    num_pools=pools_size,
+                    maxsize=maxsize,
+                    cert_reqs=cert_reqs,
+                    ca_certs=ca_certs,
+                    cert_file=configuration.cert_file,
+                    key_file=configuration.key_file,
+                    proxy_url=configuration.proxy,
+                    **addition_pool_args
+                )
+            else:
+                self.pool_manager = urllib3.ProxyManager(
+                    num_pools=pools_size,
+                    maxsize=maxsize,
+                    cert_reqs=cert_reqs,
+                    ca_certs=ca_certs,
+                    cert_file=configuration.cert_file,
+                    key_file=configuration.key_file,
+                    proxy_url=configuration.proxy,
+                    **addition_pool_args
+                )
         else:
             self.pool_manager = urllib3.PoolManager(
                 num_pools=pools_size,
@@ -140,7 +150,7 @@ def request(self, method, url, query_params=None, headers=None,
 
         timeout = None
         if _request_timeout:
-            if isinstance(_request_timeout, (int, ) if six.PY3 else (int, long)):  # noqa: E501,F821
+            if isinstance(_request_timeout, (int,) if six.PY3 else (int, long)):  # noqa: E501,F821
                 timeout = urllib3.Timeout(total=_request_timeout)
             elif (isinstance(_request_timeout, tuple) and
                   len(_request_timeout) == 2):
@@ -311,7 +321,7 @@ def __init__(self, status=None, reason=None, http_resp=None):
 
     def __str__(self):
         """Custom error messages for exception"""
-        error_message = "({0})\n"\
+        error_message = "({0})\n" \
                         "Reason: {1}\n".format(self.status, self.reason)
         if self.headers:
             error_message += "HTTP response headers: {0}\n".format(