From 769bc57ec7b0271a7cb018becee8ad156cf82704 Mon Sep 17 00:00:00 2001 From: itaru2622 Date: Wed, 29 Sep 2021 09:18:55 +0900 Subject: [PATCH] add proxy authentication supporting for websocket (stream/ws_client.py) --- stream/ws_client.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/stream/ws_client.py b/stream/ws_client.py index 356440c8a..2a60a8be9 100644 --- a/stream/ws_client.py +++ b/stream/ws_client.py @@ -29,6 +29,7 @@ from six.moves.urllib.parse import urlencode, urlparse, urlunparse from six import StringIO from websocket import WebSocket, ABNF, enableTrace +from base64 import b64decode STDIN_CHANNEL = 0 STDOUT_CHANNEL = 1 @@ -445,11 +446,20 @@ 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: 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) + 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]) }) + + websocket.connect(url, **connect_opt) return websocket