Merge branch 'master' into suppress_origin_fix

5 files changed, 41 insertions, 13 deletions

diff --git a/websocket/__init__.py b/websocket/__init__.py
index fe0c6f7..b7c6fba 100644
--- a/websocket/__init__.py
+++ b/websocket/__init__.py
@@ -26,4 +26,4 @@ from ._exceptions import *
 from ._logging import *
 from ._socket import *
 
-__version__ = "0.53.0"
+__version__ = "0.54.0"
diff --git a/websocket/_app.py b/websocket/_app.py
index 69f7e20..a05a2e3 100644
--- a/websocket/_app.py
+++ b/websocket/_app.py
@@ -48,7 +48,7 @@ class Dispatcher:
     def read(self, sock, read_callback, check_callback):
         while self.app.sock.connected:
             r, w, e = select.select(
-            (self.app.sock.sock, ), (), (), self.ping_timeout)
+                    (self.app.sock.sock, ), (), (), self.ping_timeout)
             if r:
                 if not read_callback():
                     break
@@ -118,7 +118,7 @@ class WebSocketApp(object):
           The 2nd argument is utf-8 string which we get from the server.
           The 3rd argument is data type. ABNF.OPCODE_TEXT or ABNF.OPCODE_BINARY will be came.
           The 4th argument is continue flag. if 0, the data continue
-        keep_running: this parameter is obosleted and ignored it.
+        keep_running: this parameter is obsolete and ignored.
         get_mask_key: a callable to produce new mask keys,
           see the WebSocket.set_mask_key's docstring for more information
         subprotocols: array of available sub protocols. default is None.
@@ -179,7 +179,7 @@ class WebSocketApp(object):
                     http_no_proxy=None, http_proxy_auth=None,
                     skip_utf8_validation=False,
                     host=None, origin=None, dispatcher=None,
-                    suppress_origin = False):
+                    suppress_origin = False, proxy_type=None):
         """
         run event loop for WebSocket framework.
         This loop is infinite loop and is alive during websocket is available.
@@ -199,15 +199,20 @@ class WebSocketApp(object):
         origin: update origin header.
         dispatcher: customize reading data from socket.
         suppress_origin: suppress outputting origin header.
+
+        Returns
+        -------
+        False if caught KeyboardInterrupt
+        True if other exception was raised during a loop
         """
 
-        if ping_timeout is not None and (not ping_timeout or ping_timeout <= 0):
+        if ping_timeout is not None and ping_timeout <= 0:
             ping_timeout = None
         if ping_timeout and ping_interval and ping_interval <= ping_timeout:
             raise WebSocketException("Ensure ping_interval > ping_timeout")
-        if sockopt is None:
+        if not sockopt:
             sockopt = []
-        if sslopt is None:
+        if not sslopt:
             sslopt = {}
         if self.sock:
             raise WebSocketException("socket is already opened")
@@ -232,7 +237,7 @@ class WebSocketApp(object):
         try:
             self.sock = WebSocket(
                 self.get_mask_key, sockopt=sockopt, sslopt=sslopt,
-                fire_cont_frame=self.on_cont_message and True or False,
+                fire_cont_frame=self.on_cont_message is not None,
                 skip_utf8_validation=skip_utf8_validation,
                 enable_multithread=True if ping_interval else False)
             self.sock.settimeout(getdefaulttimeout())
@@ -241,7 +246,8 @@ class WebSocketApp(object):
                 http_proxy_host=http_proxy_host,
                 http_proxy_port=http_proxy_port, http_no_proxy=http_no_proxy,
                 http_proxy_auth=http_proxy_auth, subprotocols=self.subprotocols,
-                host=host, origin=origin, suppress_origin=suppress_origin)
+                host=host, origin=origin, suppress_origin=suppress_origin,
+                proxy_type=proxy_type)
             if not dispatcher:
                 dispatcher = self.create_dispatcher(ping_timeout)
 
@@ -300,6 +306,7 @@ class WebSocketApp(object):
                 # propagate SystemExit further
                 raise
             teardown()
+            return not isinstance(e, KeyboardInterrupt)
 
     def create_dispatcher(self, ping_timeout):
         timeout = ping_timeout or 10
diff --git a/websocket/_core.py b/websocket/_core.py
index 7b76111..c91ad63 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
 
diff --git a/websocket/_http.py b/websocket/_http.py
index 8d82521..fe97480 100644
--- a/websocket/_http.py
+++ b/websocket/_http.py
@@ -101,7 +101,7 @@ def _open_proxied_socket(url, options, proxy):
 
 
 def connect(url, options, proxy, socket):
-    if proxy.host and not socket and not(proxy.type == 'http'):
+    if proxy.host and not socket and not (proxy.type == 'http'):
         return _open_proxied_socket(url, options, proxy)
 
     hostname, port, resource, is_secure = parse_url(url)