diff options
author | liris <liris.pp@gmail.com> | 2018-10-28 10:24:02 +0900 |
---|---|---|
committer | GitHub <noreply@github.com> | 2018-10-28 10:24:02 +0900 |
commit | 0153ae416db9c21b172b7a249911bd29579d978d (patch) | |
tree | 2b74bfdcd8ed7e8c49bc34452358143c0f59ca49 /websocket | |
parent | 21c6c323fa6dc6ad9bcdf2ccf3db1ea12cfe6730 (diff) | |
parent | fc09e7e2610bbf0e7308708a71eab066184aee01 (diff) | |
download | websocket-client-0153ae416db9c21b172b7a249911bd29579d978d.tar.gz |
Merge pull request #491 from partojkander/support-redirect
Diffstat (limited to 'websocket')
-rw-r--r-- | websocket/_core.py | 9 | ||||
-rw-r--r-- | websocket/_handshake.py | 18 |
2 files changed, 24 insertions, 3 deletions
diff --git a/websocket/_core.py b/websocket/_core.py index aea16e6..a85fc78 100644 --- a/websocket/_core.py +++ b/websocket/_core.py @@ -210,6 +210,7 @@ class WebSocket(object): "http_proxy_auth" - http proxy auth information. tuple of username and password. default is None + "redirect_limit" -> number of redirects to follow. "subprotocols" - array of available sub protocols. default is None. "socket" - pre-initialized stream socket. @@ -220,6 +221,13 @@ class WebSocket(object): try: self.handshake_response = handshake(self.sock, *addrs, **options) + for attempt in range(options.pop('redirect_limit', 3)): + if self.handshake_response.status in SUPPORTED_REDIRECT_STATUSES: + url = self.handshake_response.headers['location'] + self.sock.close() + self.sock, addrs = connect(url, self.sock_opt, proxy_info(**options), + options.pop('socket', None)) + self.handshake_response = handshake(self.sock, *addrs, **options) self.connected = True except: if self.sock: @@ -482,6 +490,7 @@ def create_connection(url, timeout=None, class_=WebSocket, **options): tuple of username and password. default is None "enable_multithread" -> enable lock for multithread. + "redirect_limit" -> number of redirects to follow. "sockopt" -> socket options "sslopt" -> ssl option "subprotocols" - array of available sub protocols. diff --git a/websocket/_handshake.py b/websocket/_handshake.py index d9dae47..3b8ddc0 100644 --- a/websocket/_handshake.py +++ b/websocket/_handshake.py @@ -36,7 +36,15 @@ if six.PY3: else: from base64 import encodestring as base64encode -__all__ = ["handshake_response", "handshake"] +if six.PY3: + if six.PY34: + from http import client as HTTPStatus + else: + from http import HTTPStatus +else: + import httplib as HTTPStatus + +__all__ = ["handshake_response", "handshake", "SUPPORTED_REDIRECT_STATUSES"] if hasattr(hmac, "compare_digest"): compare_digest = hmac.compare_digest @@ -47,6 +55,8 @@ else: # websocket supported version. VERSION = 13 +SUPPORTED_REDIRECT_STATUSES = [HTTPStatus.MOVED_PERMANENTLY, HTTPStatus.FOUND, HTTPStatus.SEE_OTHER] + CookieJar = SimpleCookieJar() @@ -67,6 +77,8 @@ def handshake(sock, hostname, port, resource, **options): dump("request header", header_str) status, resp = _get_resp_headers(sock) + if status in SUPPORTED_REDIRECT_STATUSES: + return handshake_response(status, resp, None) success, subproto = _validate(resp, key, options.get("subprotocols")) if not success: raise WebSocketException("Invalid WebSocket Header") @@ -134,9 +146,9 @@ def _get_handshake_headers(resource, host, port, options): return headers, key -def _get_resp_headers(sock, success_status=101): +def _get_resp_headers(sock, success_statuses=(101, 301, 302, 303)): status, resp_headers, status_message = read_headers(sock) - if status != success_status: + if status not in success_statuses: raise WebSocketBadStatusException("Handshake status %d %s", status, status_message) return status, resp_headers |