diff --git a/stream/ws_client.py b/stream/ws_client.py index 356440c..4b26ddd 100644 --- a/stream/ws_client.py +++ b/stream/ws_client.py @@ -29,6 +29,7 @@ from six import StringIO from websocket import WebSocket, ABNF, enableTrace +from base64 import urlsafe_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 = urlsafe_b64decode(b64value).decode().split(':') + connect_opt.update({'http_proxy_auth': (auth[0], auth[1]) }) + return(connect_opt) def websocket_call(configuration, _method, url, **kwargs): diff --git a/stream/ws_client_test.py b/stream/ws_client_test.py index a8f4049..bfcd64d 100644 --- a/stream/ws_client_test.py +++ b/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()