add proxy authentication supporting for websocket (stream/ws_client.py)

stream/ws_client.py

@@ -29,6 +29,7 @@
 from six import StringIO
 
 from websocket import WebSocket, ABNF, enableTrace
+from base64 import b64decode
 
 STDIN_CHANNEL = 0
 STDOUT_CHANNEL = 1
@@ -445,12 +446,27 @@ def create_websocket(configuration, url, headers=None):
         ssl_opts['keyfile'] = configuration.key_file
 
     websocket = WebSocket(sslopt=ssl_opts, skip_utf8_validation=False)
+    connect_opt = {
+         'header': header
+    }
+
+    if configuration.proxy or coniguration.proxy_headers:
+        connect_opt = websocket_proxycare(connect_opt, configuration, url, headers)
+
+    websocket.connect(url, **connect_opt)
+    return websocket
+
+def websocket_proxycare(connect_opt, configuration, url, headers):
     if configuration.proxy:
         proxy_url = urlparse(configuration.proxy)
-        websocket.connect(url, header=header, http_proxy_host=proxy_url.hostname, http_proxy_port=proxy_url.port)
-    else:
-        websocket.connect(url, header=header)
-    return websocket
+        connect_opt.update({'http_proxy_host': proxy_url.hostname, 'http_proxy_port': proxy_url.port})
+    if configuration.proxy_headers:
+        for key,value in configuration.proxy_headers.items():
+            if key == 'proxy-authorization' and value.startswith('Basic'):
+                b64value = value.split()[1]
+                auth = b64decode(b64value).decode().split(':')
+                connect_opt.update({'http_proxy_auth': (auth[0], auth[1]) })
+    return(connect_opt)
 
 
 def websocket_call(configuration, _method, url, **kwargs):

stream/ws_client_test.py

@@ -15,7 +15,21 @@
 import unittest
 
 from .ws_client import get_websocket_url
+from .ws_client import websocket_proxycare
+from kubernetes.client.configuration import Configuration
 
+try:
+    import urllib3
+    urllib3.disable_warnings()
+except ImportError:
+    pass
+
+def dictval(dict, key, default=None):
+    try:
+        val = dict[key]
+    except KeyError:
+        val = default
+    return val
 
 class WSClientTest(unittest.TestCase):
 
@@ -32,6 +46,21 @@ def test_websocket_client(self):
                 ]:
             self.assertEqual(get_websocket_url(url), ws_url)
 
+    def test_websocket_proxycare(self):
+        for proxy, idpass, expect_host, expect_port, expect_auth in [
+                ( None,                             None,        None,                None, None ),
+                ( 'http://proxy.example.com:8080/', None,        'proxy.example.com', 8080, None ),
+                ( 'http://proxy.example.com:8080/', 'user:pass', 'proxy.example.com', 8080, ('user','pass'))
+                ]:
+            config = Configuration()
+            if proxy is not None:
+                setattr(config, 'proxy', proxy)
+            if idpass is not None:
+                setattr(config, 'proxy_headers', urllib3.util.make_headers(proxy_basic_auth=idpass))
+            connect_opt = websocket_proxycare( {}, config, None, None)
+            self.assertEqual( dictval(connect_opt,'http_proxy_host'), expect_host)
+            self.assertEqual( dictval(connect_opt,'http_proxy_port'), expect_port)
+            self.assertEqual( dictval(connect_opt,'http_proxy_auth'), expect_auth)
 
 if __name__ == '__main__':
     unittest.main()