openstack/deb-python-websocket-client
Code Issues Proposed changes

fixed #125

Browse Source 
but I didn't test it, yet, because I don't have envrioment to test.
This commit is contained in:
liris
2014-11-07 09:49:45 +09:00
parent 6b029168cf
commit 19f2ef97d9
3 changed files with 66 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
ChangeLog
View File

@@ -3,8 +3,9 @@ ChangeLog
- 0.22.0 - 0.22.0
- Fix not thread-safe of Websocket.close() (#120) - Fix not thread-safe of Websocket.close() (#120)
- Try to get proxy info from environment if not explicitly provided (#124) - Try to get proxy info from environment if not explicitly provided (#124)
- support proxy basic authenticaiton. (#125)
- 0.21.0 - 0.21.0

29
websocket/_core.py
View File

@@ -172,17 +172,20 @@ def _parse_url(url):
def _get_proxy_info(is_secure, **options): def _get_proxy_info(is_secure, **options):
""" """
try to retrieve proxy host and port from environment if not provided in 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) is_secure: is the connection secure? (wss)
looks for "https_proxy" in env before falling back to "http_proxy" looks for "https_proxy" in env before falling back to "http_proxy"
options: "http_proxy_host" - http proxy host name. options: "http_proxy_host" - http proxy host name.
"http_proxy_port" - http proxy port. "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) http_proxy_host = options.get("http_proxy_host", None)
if http_proxy_host: 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"] env_keys = ["http_proxy"]
if is_secure: if is_secure:
@@ -192,9 +195,10 @@ def _get_proxy_info(is_secure, **options):
value = os.environ.get(key, None) value = os.environ.get(key, None)
if value: if value:
proxy = urlparse(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): def create_connection(url, timeout=None, **options):
@@ -220,6 +224,8 @@ def create_connection(url, timeout=None, **options):
"cookie" -> cookie value. "cookie" -> cookie value.
"http_proxy_host" - http proxy host name. "http_proxy_host" - http proxy host name.
"http_proxy_port" - http proxy port. If not set, set to 80. "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. "enable_multithread" -> enable lock for multithread.
"sockopt" -> socket options "sockopt" -> socket options
"sslopt" -> ssl option "sslopt" -> ssl option
@@ -428,12 +434,14 @@ class WebSocket(object):
"cookie" -> cookie value. "cookie" -> cookie value.
"http_proxy_host" - http proxy host name. "http_proxy_host" - http proxy host name.
"http_proxy_port" - http proxy port. If not set, set to 80. "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. "subprotocols" - array of available sub protocols. default is None.
""" """
hostname, port, resource, is_secure = _parse_url(url) 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: if not proxy_host:
addrinfo_list = socket.getaddrinfo(hostname, port, 0, 0, socket.SOL_TCP) addrinfo_list = socket.getaddrinfo(hostname, port, 0, 0, socket.SOL_TCP)
else: else:
@@ -469,7 +477,7 @@ class WebSocket(object):
raise err raise err
if proxy_host: if proxy_host:
self._tunnel(hostname, port) self._tunnel(hostname, port, proxy_auth)
if is_secure: if is_secure:
if HAVE_SSL: if HAVE_SSL:
@@ -489,9 +497,16 @@ class WebSocket(object):
self._handshake(hostname, port, resource, **options) self._handshake(hostname, port, resource, **options)
def _tunnel(self, host, port): def _tunnel(self, host, port, auth):
logger.debug("Connecting proxy...") logger.debug("Connecting proxy...")
connect_header = "CONNECT %s:%d HTTP/1.0\r\n" % (host, port) 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" connect_header += "\r\n"
_dump("request header", connect_header) _dump("request header", connect_header)

51
websocket/tests/test_websocket.py
View File

@@ -543,30 +543,61 @@ class ProxyInfoTest(unittest.TestCase):
del os.environ["https_proxy"] del os.environ["https_proxy"]
def testProxyFromArgs(self): 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"), ("localhost", 0, None))
self.assertEqual(_get_proxy_info(False, http_proxy_host="localhost", http_proxy_port=3128), ("localhost", 3128)) 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)) 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)) 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): def testProxyFromEnv(self):
os.environ["http_proxy"] = "http://localhost/" 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/" 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["http_proxy"] = "http://localhost/"
os.environ["https_proxy"] = "http://localhost2/" 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["http_proxy"] = "http://localhost:3128/"
os.environ["https_proxy"] = "http://localhost2: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["http_proxy"] = "http://localhost/"
os.environ["https_proxy"] = "http://localhost2/" 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["http_proxy"] = "http://localhost:3128/"
os.environ["https_proxy"] = "http://localhost2: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__": if __name__ == "__main__":