From 19f2ef97d925c0a0a54fc7129224dd920cec0ccb Mon Sep 17 00:00:00 2001 From: liris Date: Fri, 7 Nov 2014 09:49:45 +0900 Subject: [PATCH] fixed #125 but I didn't test it, yet, because I don't have envrioment to test. --- ChangeLog | 5 +-- websocket/_core.py | 29 +++++++++++++----- websocket/tests/test_websocket.py | 51 +++++++++++++++++++++++++------ 3 files changed, 66 insertions(+), 19 deletions(-) diff --git a/ChangeLog b/ChangeLog index 537a607..3e00dd8 100644 --- a/ChangeLog +++ b/ChangeLog @@ -3,8 +3,9 @@ ChangeLog - 0.22.0 - - Fix not thread-safe of Websocket.close() (#120) - - Try to get proxy info from environment if not explicitly provided (#124) + - Fix not thread-safe of Websocket.close() (#120) + - Try to get proxy info from environment if not explicitly provided (#124) + - support proxy basic authenticaiton. (#125) - 0.21.0 diff --git a/websocket/_core.py b/websocket/_core.py index 3b790fc..a5b56e7 100644 --- a/websocket/_core.py +++ b/websocket/_core.py @@ -172,17 +172,20 @@ def _parse_url(url): def _get_proxy_info(is_secure, **options): """ try to retrieve proxy host and port from environment if not provided in options. - result is (proxy_host, proxy_port) + result is (proxy_host, proxy_port, proxy_auth). + proxy_auth is tuple of username and password of proxy authentication information. is_secure: is the connection secure? (wss) looks for "https_proxy" in env before falling back to "http_proxy" options: "http_proxy_host" - http proxy host name. "http_proxy_port" - http proxy port. + "http_proxy_auth" - http proxy auth infomation. tuple of username and password. + defualt is None """ http_proxy_host = options.get("http_proxy_host", None) if http_proxy_host: - return http_proxy_host, options.get("http_proxy_port", 0) + return http_proxy_host, options.get("http_proxy_port", 0), options.get("http_proxy_auth", None) env_keys = ["http_proxy"] if is_secure: @@ -192,9 +195,10 @@ def _get_proxy_info(is_secure, **options): value = os.environ.get(key, None) if value: proxy = urlparse(value) - return proxy.hostname, proxy.port + auth = (proxy.username, proxy.password) if proxy.username else None + return proxy.hostname, proxy.port, auth - return None, 0 + return None, 0, None def create_connection(url, timeout=None, **options): @@ -220,6 +224,8 @@ def create_connection(url, timeout=None, **options): "cookie" -> cookie value. "http_proxy_host" - http proxy host name. "http_proxy_port" - http proxy port. If not set, set to 80. + "http_proxy_auth" - http proxy auth infomation. tuple of username and password. + defualt is None "enable_multithread" -> enable lock for multithread. "sockopt" -> socket options "sslopt" -> ssl option @@ -428,12 +434,14 @@ class WebSocket(object): "cookie" -> cookie value. "http_proxy_host" - http proxy host name. "http_proxy_port" - http proxy port. If not set, set to 80. + "http_proxy_auth" - http proxy auth infomation. tuple of username and password. + defualt is None "subprotocols" - array of available sub protocols. default is None. """ hostname, port, resource, is_secure = _parse_url(url) - proxy_host, proxy_port = _get_proxy_info(is_secure, **options) + proxy_host, proxy_port, proxy_auth = _get_proxy_info(is_secure, **options) if not proxy_host: addrinfo_list = socket.getaddrinfo(hostname, port, 0, 0, socket.SOL_TCP) else: @@ -469,7 +477,7 @@ class WebSocket(object): raise err if proxy_host: - self._tunnel(hostname, port) + self._tunnel(hostname, port, proxy_auth) if is_secure: if HAVE_SSL: @@ -489,9 +497,16 @@ class WebSocket(object): self._handshake(hostname, port, resource, **options) - def _tunnel(self, host, port): + def _tunnel(self, host, port, auth): logger.debug("Connecting proxy...") connect_header = "CONNECT %s:%d HTTP/1.0\r\n" % (host, port) + # TODO: support digest auth. + if auth and auth[0]: + auth_str = auth[0] + if auth[1]: + auth_str += ":" + auth[1] + encoded_str = base64encode(auth_str.encode()).strip().decode() + connect_header += "Proxy-Authorization: Basic %s\r\n" % encoded_str connect_header += "\r\n" _dump("request header", connect_header) diff --git a/websocket/tests/test_websocket.py b/websocket/tests/test_websocket.py index 7c29412..a6cae28 100644 --- a/websocket/tests/test_websocket.py +++ b/websocket/tests/test_websocket.py @@ -543,30 +543,61 @@ class ProxyInfoTest(unittest.TestCase): del os.environ["https_proxy"] def testProxyFromArgs(self): - self.assertEqual(_get_proxy_info(False, http_proxy_host="localhost"), ("localhost", 0)) - self.assertEqual(_get_proxy_info(False, http_proxy_host="localhost", http_proxy_port=3128), ("localhost", 3128)) - self.assertEqual(_get_proxy_info(True, http_proxy_host="localhost"), ("localhost", 0)) - self.assertEqual(_get_proxy_info(True, http_proxy_host="localhost", http_proxy_port=3128), ("localhost", 3128)) + self.assertEqual(_get_proxy_info(False, http_proxy_host="localhost"), ("localhost", 0, None)) + self.assertEqual(_get_proxy_info(False, http_proxy_host="localhost", http_proxy_port=3128), ("localhost", 3128, None)) + self.assertEqual(_get_proxy_info(True, http_proxy_host="localhost"), ("localhost", 0, None)) + self.assertEqual(_get_proxy_info(True, http_proxy_host="localhost", http_proxy_port=3128), ("localhost", 3128, None)) + + self.assertEqual(_get_proxy_info(False, http_proxy_host="localhost", http_proxy_auth=("a", "b")), + ("localhost", 0, ("a", "b"))) + self.assertEqual(_get_proxy_info(False, http_proxy_host="localhost", http_proxy_port=3128, http_proxy_auth=("a", "b")), + ("localhost", 3128, ("a", "b"))) + self.assertEqual(_get_proxy_info(True, http_proxy_host="localhost", http_proxy_auth=("a", "b")), + ("localhost", 0, ("a", "b"))) + self.assertEqual(_get_proxy_info(True, http_proxy_host="localhost", http_proxy_port=3128, http_proxy_auth=("a", "b")), + ("localhost", 3128, ("a", "b"))) + def testProxyFromEnv(self): os.environ["http_proxy"] = "http://localhost/" - self.assertEqual(_get_proxy_info(False), ("localhost", None)) + self.assertEqual(_get_proxy_info(False), ("localhost", None, None)) os.environ["http_proxy"] = "http://localhost:3128/" - self.assertEqual(_get_proxy_info(False), ("localhost", 3128)) + self.assertEqual(_get_proxy_info(False), ("localhost", 3128, None)) os.environ["http_proxy"] = "http://localhost/" os.environ["https_proxy"] = "http://localhost2/" - self.assertEqual(_get_proxy_info(False), ("localhost", None)) + self.assertEqual(_get_proxy_info(False), ("localhost", None, None)) os.environ["http_proxy"] = "http://localhost:3128/" os.environ["https_proxy"] = "http://localhost2:3128/" - self.assertEqual(_get_proxy_info(False), ("localhost", 3128)) + self.assertEqual(_get_proxy_info(False), ("localhost", 3128, None)) os.environ["http_proxy"] = "http://localhost/" os.environ["https_proxy"] = "http://localhost2/" - self.assertEqual(_get_proxy_info(True), ("localhost2", None)) + self.assertEqual(_get_proxy_info(True), ("localhost2", None, None)) os.environ["http_proxy"] = "http://localhost:3128/" os.environ["https_proxy"] = "http://localhost2:3128/" - self.assertEqual(_get_proxy_info(True), ("localhost2", 3128)) + self.assertEqual(_get_proxy_info(True), ("localhost2", 3128, None)) + + + os.environ["http_proxy"] = "http://a:b@localhost/" + self.assertEqual(_get_proxy_info(False), ("localhost", None, ("a", "b"))) + os.environ["http_proxy"] = "http://a:b@localhost:3128/" + self.assertEqual(_get_proxy_info(False), ("localhost", 3128, ("a", "b"))) + + os.environ["http_proxy"] = "http://a:b@localhost/" + os.environ["https_proxy"] = "http://a:b@localhost2/" + self.assertEqual(_get_proxy_info(False), ("localhost", None, ("a", "b"))) + os.environ["http_proxy"] = "http://a:b@localhost:3128/" + os.environ["https_proxy"] = "http://a:b@localhost2:3128/" + self.assertEqual(_get_proxy_info(False), ("localhost", 3128, ("a", "b"))) + + os.environ["http_proxy"] = "http://a:b@localhost/" + os.environ["https_proxy"] = "http://a:b@localhost2/" + self.assertEqual(_get_proxy_info(True), ("localhost2", None, ("a", "b"))) + os.environ["http_proxy"] = "http://a:b@localhost:3128/" + os.environ["https_proxy"] = "http://a:b@localhost2:3128/" + self.assertEqual(_get_proxy_info(True), ("localhost2", 3128, ("a", "b"))) + if __name__ == "__main__":