Update websocket-client to 0.57.0

lib/websocket/_handshake.py

@@ -19,33 +19,55 @@ Copyright (C) 2010 Hiroki Ohtani(liris)
     Boston, MA  02110-1335  USA
 
 """
+import hashlib
+import hmac
+import os
 
 import six
-if six.PY3:
+
+from ._cookiejar import SimpleCookieJar
+from ._exceptions import *
+from ._http import *
+from ._logging import *
+from ._socket import *
+
+if hasattr(six, 'PY3') and six.PY3:
     from base64 import encodebytes as base64encode
 else:
     from base64 import encodestring as base64encode
 
-import uuid
-import hashlib
+if hasattr(six, 'PY3') and six.PY3:
+    if hasattr(six, 'PY34') and six.PY34:
+        from http import client as HTTPStatus
+    else:
+        from http import HTTPStatus
+else:
+    import httplib as HTTPStatus
 
-from ._logging import *
-from ._url import *
-from ._socket import*
-from ._http import *
-from ._exceptions import *
+__all__ = ["handshake_response", "handshake", "SUPPORTED_REDIRECT_STATUSES"]
 
-__all__ = ["handshake_response", "handshake"]
+if hasattr(hmac, "compare_digest"):
+    compare_digest = hmac.compare_digest
+else:
+    def compare_digest(s1, s2):
+        return s1 == s2
 
 # websocket supported version.
 VERSION = 13
 
+SUPPORTED_REDIRECT_STATUSES = (HTTPStatus.MOVED_PERMANENTLY, HTTPStatus.FOUND, HTTPStatus.SEE_OTHER,)
+SUCCESS_STATUSES = SUPPORTED_REDIRECT_STATUSES + (HTTPStatus.SWITCHING_PROTOCOLS,)
+
+CookieJar = SimpleCookieJar()
+
 
 class handshake_response(object):
+
     def __init__(self, status, headers, subprotocol):
         self.status = status
         self.headers = headers
         self.subprotocol = subprotocol
+        CookieJar.add(headers.get("set-cookie"))
 
 
 def handshake(sock, hostname, port, resource, **options):
@@ -56,6 +78,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")
@@ -63,38 +87,68 @@ def handshake(sock, hostname, port, resource, **options):
     return handshake_response(status, resp, subproto)
 
 
-def _get_handshake_headers(resource, host, port, options):
-    headers = []
-    headers.append("GET %s HTTP/1.1" % resource)
-    headers.append("Upgrade: websocket")
-    headers.append("Connection: Upgrade")
-    if port == 80:
-        hostport = host
-    else:
-        hostport = "%s:%d" % (host, port)
+def _pack_hostname(hostname):
+    # IPv6 address
+    if ':' in hostname:
+        return '[' + hostname + ']'
 
-    if "host" in options and options["host"]:
+    return hostname
+
+def _get_handshake_headers(resource, host, port, options):
+    headers = [
+        "GET %s HTTP/1.1" % resource,
+        "Upgrade: websocket"
+    ]
+    if port == 80 or port == 443:
+        hostport = _pack_hostname(host)
+    else:
+        hostport = "%s:%d" % (_pack_hostname(host), port)
+    if "host" in options and options["host"] is not None:
         headers.append("Host: %s" % options["host"])
     else:
         headers.append("Host: %s" % hostport)
 
-    if "origin" in options and options["origin"]:
-        headers.append("Origin: %s" % options["origin"])
-    else:
-        headers.append("Origin: http://%s" % hostport)
+    if "suppress_origin" not in options or not options["suppress_origin"]:
+        if "origin" in options and options["origin"] is not None:
+            headers.append("Origin: %s" % options["origin"])
+        else:
+            headers.append("Origin: http://%s" % hostport)
 
     key = _create_sec_websocket_key()
-    headers.append("Sec-WebSocket-Key: %s" % key)
-    headers.append("Sec-WebSocket-Version: %s" % VERSION)
+    
+    # Append Sec-WebSocket-Key & Sec-WebSocket-Version if not manually specified
+    if not 'header' in options or 'Sec-WebSocket-Key' not in options['header']:
+        key = _create_sec_websocket_key()
+        headers.append("Sec-WebSocket-Key: %s" % key)
+    else:
+        key = options['header']['Sec-WebSocket-Key']
+
+    if not 'header' in options or 'Sec-WebSocket-Version' not in options['header']:
+        headers.append("Sec-WebSocket-Version: %s" % VERSION)
+
+    if not 'connection' in options or options['connection'] is None:
+        headers.append('Connection: upgrade')
+    else:
+        headers.append(options['connection'])
 
     subprotocols = options.get("subprotocols")
     if subprotocols:
         headers.append("Sec-WebSocket-Protocol: %s" % ",".join(subprotocols))
 
     if "header" in options:
-        headers.extend(options["header"])
+        header = options["header"]
+        if isinstance(header, dict):
+            header = [
+                ": ".join([k, v])
+                for k, v in header.items()
+                if v is not None
+            ]
+        headers.extend(header)
 
-    cookie = options.get("cookie", None)
+    server_cookie = CookieJar.get(host)
+    client_cookie = options.get("cookie", None)
+
+    cookie = "; ".join(filter(None, [server_cookie, client_cookie]))
 
     if cookie:
         headers.append("Cookie: %s" % cookie)
@@ -105,16 +159,17 @@ def _get_handshake_headers(resource, host, port, options):
     return headers, key
 
 
-def _get_resp_headers(sock, success_status=101):
-    status, resp_headers = read_headers(sock)
-    if status != success_status:
-        raise WebSocketException("Handshake status %d" % status)
+def _get_resp_headers(sock, success_statuses=SUCCESS_STATUSES):
+    status, resp_headers, status_message = read_headers(sock)
+    if status not in success_statuses:
+        raise WebSocketBadStatusException("Handshake status %d %s", status, status_message, resp_headers)
     return status, resp_headers
 
+
 _HEADERS_TO_CHECK = {
     "upgrade": "websocket",
     "connection": "upgrade",
-    }
+}
 
 
 def _validate(headers, key, subprotocols):
@@ -143,7 +198,8 @@ def _validate(headers, key, subprotocols):
 
     value = (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode('utf-8')
     hashed = base64encode(hashlib.sha1(value).digest()).strip().lower()
-    success = (hashed == result)
+    success = compare_digest(hashed, result)
+
     if success:
         return True, subproto
     else:
@@ -151,5 +207,5 @@ def _validate(headers, key, subprotocols):
 
 
 def _create_sec_websocket_key():
-    uid = uuid.uuid4()
-    return base64encode(uid.bytes).decode('utf-8').strip()
+    randomness = os.urandom(16)
+    return base64encode(randomness).decode('utf-8').strip()