From 877727110956253be05e45dfb0e18bd094c54e90 Mon Sep 17 00:00:00 2001
From: itaru2622 <itaru2622@gmail.com>
Date: Tue, 5 Oct 2021 20:50:01 +0900
Subject: [PATCH] proxy authentication supporting for websocket
 (stream/ws_client.py), with unittest

---
 stream/ws_client.py      | 22 ++++++++++++++--------
 stream/ws_client_test.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/stream/ws_client.py b/stream/ws_client.py
index 2a60a8be9..419d28b20 100644
--- a/stream/ws_client.py
+++ b/stream/ws_client.py
@@ -449,18 +449,24 @@ def create_websocket(configuration, url, headers=None):
     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)
         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
+        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):
diff --git a/stream/ws_client_test.py b/stream/ws_client_test.py
index a8f4049dd..bfcd64d53 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 @@ class WSClientTest(unittest.TestCase):
                 ]:
             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()