From e41e5cbfdd7b778e6f68e86658e95f9e413133cb Mon Sep 17 00:00:00 2001
From: Chandler Abraham <cabraham@twitter.com>
Date: Thu, 9 Apr 2015 19:35:40 -0700
Subject: [PATCH 1/5] netlib websockets

---
 netlib/http.py                       |  14 +
 netlib/utils.py                      |   3 +
 netlib/websockets/__init__.py        |   1 +
 netlib/websockets/implementations.py |  81 ++++++
 netlib/websockets/websockets.py      | 368 +++++++++++++++++++++++++++
 test/test_websockets.py              |  15 ++
 6 files changed, 482 insertions(+)
 create mode 100644 netlib/websockets/__init__.py
 create mode 100644 netlib/websockets/implementations.py
 create mode 100644 netlib/websockets/websockets.py
 create mode 100644 test/test_websockets.py

diff --git a/netlib/http.py b/netlib/http.py
index 264388636..2c72621dd 100644
--- a/netlib/http.py
+++ b/netlib/http.py
@@ -29,6 +29,20 @@ def _is_valid_host(host):
         return None
     return True
 
+def is_successful_upgrade(request, response):
+    """
+        determines if a client and server successfully agreed to an HTTP protocol upgrade
+
+        https://developer.mozilla.org/en-US/docs/Web/HTTP/Protocol_upgrade_mechanism
+    """
+    http_switching_protocols_code = 101
+    
+    if request and response:
+        responseUpgrade = request.headers.get("Upgrade")
+        requestUpgrade  = response.headers.get("Upgrade")
+        if response.code == http_switching_protocols_code and responseUpgrade == requestUpgrade:
+            return requestUpgrade[0] if len(requestUpgrade) > 0 else None
+    return None
 
 def parse_url(url):
     """
diff --git a/netlib/utils.py b/netlib/utils.py
index 79077ac60..03a70977c 100644
--- a/netlib/utils.py
+++ b/netlib/utils.py
@@ -8,6 +8,9 @@ def isascii(s):
         return False
     return True
 
+# best way to do it in python 2.x
+def bytes_to_int(i):
+  return int(i.encode('hex'), 16)
 
 def cleanBin(s, fixspacing=False):
     """
diff --git a/netlib/websockets/__init__.py b/netlib/websockets/__init__.py
new file mode 100644
index 000000000..9b4faa337
--- /dev/null
+++ b/netlib/websockets/__init__.py
@@ -0,0 +1 @@
+from __future__ import (absolute_import, print_function, division)
diff --git a/netlib/websockets/implementations.py b/netlib/websockets/implementations.py
new file mode 100644
index 000000000..78ae5be6b
--- /dev/null
+++ b/netlib/websockets/implementations.py
@@ -0,0 +1,81 @@
+from netlib import tcp
+from base64 import b64encode
+from StringIO import StringIO
+from . import websockets as ws
+import struct
+import SocketServer
+import os
+
+# Simple websocket client and servers that are used to exercise the functionality in websockets.py
+# These are *not* fully RFC6455 compliant
+
+class WebSocketsEchoHandler(tcp.BaseHandler): 
+    def __init__(self, connection, address, server):
+        super(WebSocketsEchoHandler, self).__init__(connection, address, server)
+        self.handshake_done = False
+
+    def handle(self):
+       while True:
+          if not self.handshake_done:
+              self.handshake()
+          else:
+              self.read_next_message()
+
+    def read_next_message(self):
+        decoded = ws.WebSocketsFrame.from_byte_stream(self.rfile.read).decoded_payload
+        self.on_message(decoded)
+
+    def send_message(self, message):
+        frame = ws.WebSocketsFrame.default_frame_from_message(message, from_client = False)
+        self.wfile.write(frame.to_bytes())
+        self.wfile.flush()
+ 
+    def handshake(self):
+        client_hs = ws.read_handshake(self.rfile.read, 1)
+        key       = ws.server_process_handshake(client_hs)
+        response  = ws.create_server_handshake(key)
+        self.wfile.write(response)
+        self.wfile.flush()
+        self.handshake_done = True
+
+    def on_message(self, message):
+        if message is not None:
+            self.send_message(message)
+
+
+class WebSocketsClient(tcp.TCPClient):
+    def __init__(self, address, source_address=None):
+        super(WebSocketsClient, self).__init__(address, source_address)
+        self.version    = "13"
+        self.key        = b64encode(os.urandom(16)).decode('utf-8')
+        self.resource   = "/"
+
+    def connect(self):
+        super(WebSocketsClient, self).connect()
+
+        handshake = ws.create_client_handshake(
+            self.address.host,
+            self.address.port,
+            self.key,
+            self.version,
+            self.resource
+        )
+
+        self.wfile.write(handshake)
+        self.wfile.flush()
+
+        response = ws.read_handshake(self.rfile.read, 1)
+        
+        if not response:
+            self.close()
+
+    def read_next_message(self):
+        try:
+            return ws.WebSocketsFrame.from_byte_stream(self.rfile.read).payload
+        except IndexError:
+            self.close()
+ 
+    def send_message(self, message):
+        frame = ws.WebSocketsFrame.default_frame_from_message(message, from_client = True)
+        self.wfile.write(frame.to_bytes())
+        self.wfile.flush()
diff --git a/netlib/websockets/websockets.py b/netlib/websockets/websockets.py
new file mode 100644
index 000000000..b796ce399
--- /dev/null
+++ b/netlib/websockets/websockets.py
@@ -0,0 +1,368 @@
+from __future__ import absolute_import
+
+from base64 import b64encode
+from hashlib import sha1
+from mimetools import Message
+from netlib import tcp
+from netlib import utils
+from StringIO import StringIO
+import os
+import SocketServer
+import struct
+import io
+
+# Colleciton of utility functions that implement small portions of the RFC6455 WebSockets Protocol
+# Useful for building WebSocket clients and servers. 
+#
+# Emphassis is on readabilty, simplicity and modularity, not performance or completeness
+#
+# This is a work in progress and does not yet contain all the utilites need to create fully complient client/servers
+#
+# Spec: https://tools.ietf.org/html/rfc6455
+
+# The magic sha that websocket servers must know to prove they understand RFC6455
+websockets_magic = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
+
+class WebSocketFrameValidationException(Exception):
+    pass
+
+class WebSocketsFrame(object):
+    """
+        Represents one websockets frame.
+        Constructor takes human readable forms of the frame components
+        from_bytes() is also avaliable.
+
+        WebSockets Frame as defined in RFC6455
+       
+          0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+         +-+-+-+-+-------+-+-------------+-------------------------------+
+         |F|R|R|R| opcode|M| Payload len |    Extended payload length    |
+         |I|S|S|S|  (4)  |A|     (7)     |             (16/64)           |
+         |N|V|V|V|       |S|             |   (if payload len==126/127)   |
+         | |1|2|3|       |K|             |                               |
+         +-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +
+         |     Extended payload length continued, if payload len == 127  |
+         + - - - - - - - - - - - - - - - +-------------------------------+
+         |                               |Masking-key, if MASK set to 1  |
+         +-------------------------------+-------------------------------+
+         | Masking-key (continued)       |          Payload Data         |
+         +-------------------------------- - - - - - - - - - - - - - - - +
+         :                     Payload Data continued ...                :
+         + - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - +
+         |                     Payload Data continued ...                |
+         +---------------------------------------------------------------+
+    """
+    def __init__(
+        self,
+        fin,                          # decmial integer 1 or 0
+        opcode,                       # decmial integer 1 - 4
+        mask_bit,                     # decimal integer 1 or 0
+        payload_length_code,          # decimal integer 1 - 127
+        decoded_payload,              # bytestring
+        rsv1                  = 0,    # decimal integer 1 or 0
+        rsv2                  = 0,    # decimal integer 1 or 0
+        rsv3                  = 0,    # decimal integer 1 or 0
+        payload               = None, # bytestring 
+        masking_key           = None, # 32 bit byte string
+        actual_payload_length = None, # any decimal integer
+        use_validation        = True  # indicates whether or not you care if this frame adheres to the spec
+    ):
+        self.fin                   = fin
+        self.rsv1                  = rsv1
+        self.rsv2                  = rsv2
+        self.rsv3                  = rsv3
+        self.opcode                = opcode
+        self.mask_bit              = mask_bit
+        self.payload_length_code   = payload_length_code
+        self.masking_key           = masking_key
+        self.payload               = payload
+        self.decoded_payload       = decoded_payload
+        self.actual_payload_length = actual_payload_length
+        self.use_validation        = use_validation
+
+        if self.use_validation:
+            self.validate_frame()
+
+    @classmethod
+    def from_bytes(cls, bytestring):
+        """
+          Construct a websocket frame from an in-memory bytestring
+          to construct a frame from a stream of bytes, use read_frame() directly
+        """ 
+        self.from_byte_stream(io.BytesIO(bytestring).read)
+
+    @classmethod
+    def default_frame_from_message(cls, message, from_client = False):
+        """
+          Construct a basic websocket frame from some default values. 
+          Creates a non-fragmented text frame.
+        """ 
+        length_code, actual_length = get_payload_length_pair(message)
+
+        if from_client:
+            mask_bit    = 1
+            masking_key = random_masking_key()
+            payload     = apply_mask(message, masking_key)
+        else:
+            mask_bit    = 0
+            masking_key = None
+            payload     = message
+        
+        return cls(
+            fin                   = 1, # final frame
+            opcode                = 1, # text
+            mask_bit              = mask_bit,
+            payload_length_code   = length_code,
+            payload               = payload,
+            masking_key           = masking_key,
+            decoded_payload       = message,
+            actual_payload_length = actual_length
+        )
+
+    def validate_frame(self):
+        """
+         Validate websocket frame invariants, call at anytime to ensure the WebSocketsFrame
+         has not been corrupted.
+        """ 
+        try: 
+            assert 0 <= self.fin <= 1
+            assert 0 <= self.rsv1 <= 1
+            assert 0 <= self.rsv2 <= 1
+            assert 0 <= self.rsv3 <= 1
+            assert 1 <= self.opcode <= 4
+            assert 0 <= self.mask_bit <= 1
+            assert 1 <= self.payload_length_code <= 127
+            
+            if self.mask_bit == 1:
+                assert 1 <= len(self.masking_key) <= 4
+            else:
+                assert self.masking_key == None
+            
+            assert self.actual_payload_length == len(self.payload)
+
+            if self.payload is not None and self.masking_key is not None:
+                apply_mask(self.payload, self.masking_key) == self.decoded_payload
+
+        except AssertionError:
+            raise WebSocketFrameValidationException()
+
+    def human_readable(self):
+        return "\n".join([
+          ("fin                   - " + str(self.fin)),
+          ("rsv1                  - " + str(self.rsv1)),
+          ("rsv2                  - " + str(self.rsv2)),
+          ("rsv3                  - " + str(self.rsv3)),
+          ("opcode                - " + str(self.opcode)),
+          ("mask_bit              - " + str(self.mask_bit)),
+          ("payload_length_code   - " + str(self.payload_length_code)),
+          ("masking_key           - " + str(self.masking_key)),
+          ("payload               - " + str(self.payload)),
+          ("decoded_payload       - " + str(self.decoded_payload)),
+          ("actual_payload_length - " + str(self.actual_payload_length)),
+          ("use_validation        - " + str(self.use_validation))])
+
+    def to_bytes(self):
+        """
+          Serialize the frame back into the wire format, returns a bytestring
+        """ 
+        # validate enforces all the assumptions made by this serializer
+        # in the spritit of mitmproxy, it's possible to create and serialize invalid frames 
+        # by skipping validation.
+        if self.use_validation:
+            self.validate_frame()
+
+        max_16_bit_int = (1 << 16)
+        max_64_bit_int = (1 << 63)
+
+        # break down of the bit-math used to construct the first byte from the frame's integer values
+        # first shift the significant bit into the correct position  
+        # 00000001 << 7 = 10000000
+        # ...
+        # then combine:
+        # 
+        # 10000000 fin
+        # 01000000 res1
+        # 00100000 res2
+        # 00010000 res3
+        # 00000001 opcode
+        # -------- OR 
+        # 11110001   = first_byte
+
+        first_byte  = (self.fin << 7) | (self.rsv1 << 6) | (self.rsv2 << 4) | (self.rsv3 << 4) | self.opcode
+        
+        second_byte = (self.mask_bit << 7) | self.payload_length_code
+
+        bytes = chr(first_byte) + chr(second_byte)
+
+        if self.actual_payload_length < 126:
+            pass
+        
+        elif self.actual_payload_length < max_16_bit_int:
+            # '!H' pack as 16 bit unsigned short
+            bytes += struct.pack('!H', self.actual_payload_length) # add 2 byte extended payload length
+        
+        elif self.actual_payload_length < max_64_bit_int:
+            # '!Q' = pack as 64 bit unsigned long long
+            bytes += struct.pack('!Q', self.actual_payload_length) # add 8 bytes extended payload length
+
+        if self.masking_key is not None:
+            bytes += self.masking_key
+
+        bytes += self.payload # already will be encoded if neccessary
+
+        return bytes
+
+
+    @classmethod
+    def from_byte_stream(cls, read_bytes):
+        """
+          read a websockets frame sent by a server or client
+      
+          read_bytes is a function that can be backed
+          by sockets or by any byte reader. So this 
+          function may be used to read frames from disk/wire/memory
+        """ 
+        first_byte  = utils.bytes_to_int(read_bytes(1))
+        second_byte = utils.bytes_to_int(read_bytes(1))
+        
+        fin            = first_byte  >> 7  # grab the left most bit
+        opcode         = first_byte  & 15  # grab right most 4 bits by and-ing with 00001111
+        mask_bit       = second_byte >> 7  # grab left most bit
+        payload_length = second_byte & 127 # grab the next 7 bits
+
+        # payload_lengthy > 125 indicates you need to read more bytes
+        # to get the actual payload length
+        if payload_length <= 125:
+           actual_payload_length = payload_length
+
+        elif payload_length == 126:
+           actual_payload_length = utils.bytes_to_int(read_bytes(2))
+
+        elif payload_length == 127: 
+           actual_payload_length = utils.bytes_to_int(read_bytes(8))
+
+        # masking key only present if mask bit set
+        if mask_bit == 1:
+            masking_key = read_bytes(4)
+        else:
+            masking_key = None
+        
+        payload = read_bytes(actual_payload_length)
+        
+        if mask_bit == 1:
+            decoded_payload = apply_mask(payload, masking_key)
+        else:
+            decoded_payload = payload
+
+        return cls(
+            fin                   = fin,
+            opcode                = opcode,
+            mask_bit              = mask_bit,
+            payload_length_code   = payload_length,
+            payload               = payload,
+            masking_key           = masking_key,
+            decoded_payload       = decoded_payload,
+            actual_payload_length = actual_payload_length
+        )
+
+def apply_mask(message, masking_key):
+    """
+    Data sent from the server must be masked to prevent malicious clients
+    from sending data over the wire in predictable patterns
+
+    This method both encodes and decodes strings with the provided mask
+
+    Servers do not have to mask data they send to the client.
+    https://tools.ietf.org/html/rfc6455#section-5.3
+    """
+    masks = [utils.bytes_to_int(byte) for byte in masking_key]
+    result = ""
+    for char in message:
+        result += chr(ord(char) ^ masks[len(result) % 4])
+    return result
+
+def random_masking_key():
+    return os.urandom(4)
+
+def masking_key_list(masking_key):
+    return [utils.bytes_to_int(byte) for byte in masking_key]
+
+def create_client_handshake(host, port, key, version, resource):
+    """
+      WebSockets connections are intiated by the client with a valid HTTP upgrade request
+    """
+    headers = [
+        ('Host', '%s:%s' % (host, port)),
+        ('Connection', 'Upgrade'),
+        ('Upgrade', 'websocket'),
+        ('Sec-WebSocket-Key', key),
+        ('Sec-WebSocket-Version', version)
+    ]
+    request = "GET %s HTTP/1.1" % resource
+    return build_handshake(headers, request)
+
+
+def create_server_handshake(key, magic = websockets_magic):
+    """
+      The server response is a valid HTTP 101 response. 
+    """ 
+    digest = b64encode(sha1(key + magic).hexdigest().decode('hex'))
+    headers = [
+        ('Connection', 'Upgrade'),
+        ('Upgrade', 'websocket'),
+        ('Sec-WebSocket-Accept', digest)
+    ]
+    request = "HTTP/1.1 101 Switching Protocols"
+    return build_handshake(headers, request)
+
+
+def build_handshake(headers, request):
+    handshake = [request.encode('utf-8')]
+    for header, value in headers:
+        handshake.append(("%s: %s" % (header, value)).encode('utf-8'))
+    handshake.append(b'\r\n')
+    return b'\r\n'.join(handshake)
+
+
+def read_handshake(read_bytes, num_bytes_per_read):
+    """
+      From provided function that reads bytes, read in a 
+      complete HTTP request, which terminates with a CLRF
+    """ 
+    response   = b''
+    doubleCLRF = b'\r\n\r\n'
+    while True:
+        bytes = read_bytes(num_bytes_per_read)
+        if not bytes:
+            break
+        response += bytes
+        if doubleCLRF in response:
+            break
+    return response
+
+def get_payload_length_pair(payload_bytestring):
+    """
+     A websockets frame contains an initial length_code, and an optional
+     extended length code to represent the actual length if length code is larger
+     than 125
+    """ 
+    actual_length = len(payload_bytestring)
+  
+    if actual_length <= 125:
+        length_code = actual_length
+    elif actual_length >= 126 and actual_length <= 65535:
+        length_code = 126
+    else:
+        length_code = 127
+    return (length_code, actual_length)
+
+def server_process_handshake(handshake):
+    headers = Message(StringIO(handshake.split('\r\n', 1)[1]))
+    if headers.get("Upgrade", None) != "websocket":
+        return
+    key = headers['Sec-WebSocket-Key']
+    return key
+
+def generate_client_nounce():
+    return b64encode(os.urandom(16)).decode('utf-8')
+
diff --git a/test/test_websockets.py b/test/test_websockets.py
new file mode 100644
index 000000000..d7e1627fa
--- /dev/null
+++ b/test/test_websockets.py
@@ -0,0 +1,15 @@
+from netlib import test
+from netlib.websockets import implementations as ws
+
+class TestWebSockets(test.ServerTestBase):
+    handler = ws.WebSocketsEchoHandler
+
+    def test_websockets_echo(self):
+        msg    = "hello I'm the client"
+        client = ws.WebSocketsClient(("127.0.0.1", self.port))
+        client.connect()
+        client.send_message(msg)
+        response = client.read_next_message()
+        print "Assert response: " + response + " == msg: " + msg
+        assert response == msg
+

From 0edc04814e3affa71025938ac354707b9b4c481c Mon Sep 17 00:00:00 2001
From: Chandler Abraham <cabraham@twitter.com>
Date: Sat, 11 Apr 2015 11:35:15 -0700
Subject: [PATCH 2/5] small cleanups, working on tests

---
 netlib/websockets/implementations.py | 10 ++++----
 netlib/websockets/websockets.py      | 35 ++++++++++++++--------------
 test/test_websockets.py              | 24 +++++++++++++++----
 3 files changed, 41 insertions(+), 28 deletions(-)

diff --git a/netlib/websockets/implementations.py b/netlib/websockets/implementations.py
index 78ae5be6b..ff42ff658 100644
--- a/netlib/websockets/implementations.py
+++ b/netlib/websockets/implementations.py
@@ -26,8 +26,8 @@ class WebSocketsEchoHandler(tcp.BaseHandler):
         self.on_message(decoded)
 
     def send_message(self, message):
-        frame = ws.WebSocketsFrame.default_frame_from_message(message, from_client = False)
-        self.wfile.write(frame.to_bytes())
+        frame = ws.WebSocketsFrame.default(message, from_client = False)
+        self.wfile.write(frame.safe_to_bytes())
         self.wfile.flush()
  
     def handshake(self):
@@ -47,7 +47,7 @@ class WebSocketsClient(tcp.TCPClient):
     def __init__(self, address, source_address=None):
         super(WebSocketsClient, self).__init__(address, source_address)
         self.version    = "13"
-        self.key        = b64encode(os.urandom(16)).decode('utf-8')
+        self.key        = ws.generate_client_nounce()
         self.resource   = "/"
 
     def connect(self):
@@ -76,6 +76,6 @@ class WebSocketsClient(tcp.TCPClient):
             self.close()
  
     def send_message(self, message):
-        frame = ws.WebSocketsFrame.default_frame_from_message(message, from_client = True)
-        self.wfile.write(frame.to_bytes())
+        frame = ws.WebSocketsFrame.default(message, from_client = True)
+        self.wfile.write(frame.safe_to_bytes())
         self.wfile.flush()
diff --git a/netlib/websockets/websockets.py b/netlib/websockets/websockets.py
index b796ce399..527d55d62 100644
--- a/netlib/websockets/websockets.py
+++ b/netlib/websockets/websockets.py
@@ -65,7 +65,6 @@ class WebSocketsFrame(object):
         payload               = None, # bytestring 
         masking_key           = None, # 32 bit byte string
         actual_payload_length = None, # any decimal integer
-        use_validation        = True  # indicates whether or not you care if this frame adheres to the spec
     ):
         self.fin                   = fin
         self.rsv1                  = rsv1
@@ -78,21 +77,18 @@ class WebSocketsFrame(object):
         self.payload               = payload
         self.decoded_payload       = decoded_payload
         self.actual_payload_length = actual_payload_length
-        self.use_validation        = use_validation
-
-        if self.use_validation:
-            self.validate_frame()
 
     @classmethod
     def from_bytes(cls, bytestring):
         """
           Construct a websocket frame from an in-memory bytestring
-          to construct a frame from a stream of bytes, use read_frame() directly
+          to construct a frame from a stream of bytes, use from_byte_stream() directly
         """ 
         self.from_byte_stream(io.BytesIO(bytestring).read)
 
+
     @classmethod
-    def default_frame_from_message(cls, message, from_client = False):
+    def default(cls, message, from_client = False):
         """
           Construct a basic websocket frame from some default values. 
           Creates a non-fragmented text frame.
@@ -119,7 +115,7 @@ class WebSocketsFrame(object):
             actual_payload_length = actual_length
         )
 
-    def validate_frame(self):
+    def frame_is_valid(self):
         """
          Validate websocket frame invariants, call at anytime to ensure the WebSocketsFrame
          has not been corrupted.
@@ -141,10 +137,11 @@ class WebSocketsFrame(object):
             assert self.actual_payload_length == len(self.payload)
 
             if self.payload is not None and self.masking_key is not None:
-                apply_mask(self.payload, self.masking_key) == self.decoded_payload
+                assert apply_mask(self.payload, self.masking_key) == self.decoded_payload
 
+            return True 
         except AssertionError:
-            raise WebSocketFrameValidationException()
+            return False
 
     def human_readable(self):
         return "\n".join([
@@ -161,15 +158,19 @@ class WebSocketsFrame(object):
           ("actual_payload_length - " + str(self.actual_payload_length)),
           ("use_validation        - " + str(self.use_validation))])
 
+    def safe_to_bytes(self):
+      try:
+          assert self.frame_is_valid()
+          return self.to_bytes()
+      except:
+          raise WebSocketFrameValidationException()
+
     def to_bytes(self):
         """
           Serialize the frame back into the wire format, returns a bytestring
+          If you haven't checked is_valid_frame() then there's no guarentees that the
+          serialized bytes will be correct. see safe_to_bytes()
         """ 
-        # validate enforces all the assumptions made by this serializer
-        # in the spritit of mitmproxy, it's possible to create and serialize invalid frames 
-        # by skipping validation.
-        if self.use_validation:
-            self.validate_frame()
 
         max_16_bit_int = (1 << 16)
         max_64_bit_int = (1 << 63)
@@ -198,6 +199,7 @@ class WebSocketsFrame(object):
             pass
         
         elif self.actual_payload_length < max_16_bit_int:
+
             # '!H' pack as 16 bit unsigned short
             bytes += struct.pack('!H', self.actual_payload_length) # add 2 byte extended payload length
         
@@ -284,9 +286,6 @@ def apply_mask(message, masking_key):
 def random_masking_key():
     return os.urandom(4)
 
-def masking_key_list(masking_key):
-    return [utils.bytes_to_int(byte) for byte in masking_key]
-
 def create_client_handshake(host, port, key, version, resource):
     """
       WebSockets connections are intiated by the client with a valid HTTP upgrade request
diff --git a/test/test_websockets.py b/test/test_websockets.py
index d7e1627fa..0b2647ef3 100644
--- a/test/test_websockets.py
+++ b/test/test_websockets.py
@@ -1,15 +1,29 @@
 from netlib import test
-from netlib.websockets import implementations as ws
+from netlib.websockets import implementations as impl
+from netlib.websockets import websockets as ws
+import os
 
 class TestWebSockets(test.ServerTestBase):
-    handler = ws.WebSocketsEchoHandler
+    handler = impl.WebSocketsEchoHandler
 
-    def test_websockets_echo(self):
-        msg    = "hello I'm the client"
-        client = ws.WebSocketsClient(("127.0.0.1", self.port))
+    def echo(self, msg):
+        client = impl.WebSocketsClient(("127.0.0.1", self.port))
         client.connect()
         client.send_message(msg)
         response = client.read_next_message()
         print "Assert response: " + response + " == msg: " + msg
         assert response == msg
 
+    def test_simple_echo(self):
+        self.echo("hello I'm the client")
+
+    def test_frame_sizes(self):
+        small_string     =  os.urandom(100)   # length can fit in the the 7 bit payload length 
+        medium_string    =  os.urandom(50000) # 50kb, sligthly larger than can fit in a 7 bit int
+        large_string     =  os.urandom(150000) # 150kb, slightly larger than can fit in a 16 bit int
+
+        self.echo(small_string)
+        self.echo(medium_string)
+        self.echo(large_string)
+
+

From f131f9b855e77554072415c925ed112ec74ee48a Mon Sep 17 00:00:00 2001
From: Chandler Abraham <cabraham@twitter.com>
Date: Sat, 11 Apr 2015 15:40:18 -0700
Subject: [PATCH 3/5] handshake tests, serialization test

---
 netlib/websockets/implementations.py | 19 +++++----
 netlib/websockets/websockets.py      | 51 +++++++++++++++-------
 test/test_websockets.py              | 63 ++++++++++++++++++++++++----
 3 files changed, 105 insertions(+), 28 deletions(-)

diff --git a/netlib/websockets/implementations.py b/netlib/websockets/implementations.py
index ff42ff658..73a846905 100644
--- a/netlib/websockets/implementations.py
+++ b/netlib/websockets/implementations.py
@@ -32,7 +32,7 @@ class WebSocketsEchoHandler(tcp.BaseHandler):
  
     def handshake(self):
         client_hs = ws.read_handshake(self.rfile.read, 1)
-        key       = ws.server_process_handshake(client_hs)
+        key       = ws.process_handshake_from_client(client_hs)
         response  = ws.create_server_handshake(key)
         self.wfile.write(response)
         self.wfile.flush()
@@ -46,9 +46,9 @@ class WebSocketsEchoHandler(tcp.BaseHandler):
 class WebSocketsClient(tcp.TCPClient):
     def __init__(self, address, source_address=None):
         super(WebSocketsClient, self).__init__(address, source_address)
-        self.version    = "13"
-        self.key        = ws.generate_client_nounce()
-        self.resource   = "/"
+        self.version       = "13"
+        self.client_nounce = ws.create_client_nounce()
+        self.resource      = "/"
 
     def connect(self):
         super(WebSocketsClient, self).connect()
@@ -56,7 +56,7 @@ class WebSocketsClient(tcp.TCPClient):
         handshake = ws.create_client_handshake(
             self.address.host,
             self.address.port,
-            self.key,
+            self.client_nounce,
             self.version,
             self.resource
         )
@@ -64,9 +64,14 @@ class WebSocketsClient(tcp.TCPClient):
         self.wfile.write(handshake)
         self.wfile.flush()
 
-        response = ws.read_handshake(self.rfile.read, 1)
+        server_handshake = ws.read_handshake(self.rfile.read, 1)
         
-        if not response:
+        if not server_handshake:
+            self.close()
+
+        server_nounce = ws.process_handshake_from_server(server_handshake, self.client_nounce)
+
+        if not server_nounce == ws.create_server_nounce(self.client_nounce):
             self.close()
 
     def read_next_message(self):
diff --git a/netlib/websockets/websockets.py b/netlib/websockets/websockets.py
index 527d55d62..cf9a68aa9 100644
--- a/netlib/websockets/websockets.py
+++ b/netlib/websockets/websockets.py
@@ -84,7 +84,7 @@ class WebSocketsFrame(object):
           Construct a websocket frame from an in-memory bytestring
           to construct a frame from a stream of bytes, use from_byte_stream() directly
         """ 
-        self.from_byte_stream(io.BytesIO(bytestring).read)
+        return cls.from_byte_stream(io.BytesIO(bytestring).read)
 
 
     @classmethod
@@ -115,7 +115,7 @@ class WebSocketsFrame(object):
             actual_payload_length = actual_length
         )
 
-    def frame_is_valid(self):
+    def is_valid(self):
         """
          Validate websocket frame invariants, call at anytime to ensure the WebSocketsFrame
          has not been corrupted.
@@ -155,12 +155,11 @@ class WebSocketsFrame(object):
           ("masking_key           - " + str(self.masking_key)),
           ("payload               - " + str(self.payload)),
           ("decoded_payload       - " + str(self.decoded_payload)),
-          ("actual_payload_length - " + str(self.actual_payload_length)),
-          ("use_validation        - " + str(self.use_validation))])
+          ("actual_payload_length - " + str(self.actual_payload_length))])
 
     def safe_to_bytes(self):
       try:
-          assert self.frame_is_valid()
+          assert self.is_valid()
           return self.to_bytes()
       except:
           raise WebSocketFrameValidationException()
@@ -197,7 +196,7 @@ class WebSocketsFrame(object):
 
         if self.actual_payload_length < 126:
             pass
-        
+
         elif self.actual_payload_length < max_16_bit_int:
 
             # '!H' pack as 16 bit unsigned short
@@ -267,6 +266,20 @@ class WebSocketsFrame(object):
             actual_payload_length = actual_payload_length
         )
 
+    def __eq__(self, other):
+        return (
+         self.fin                   == other.fin and
+         self.rsv1                  == other.rsv1 and
+         self.rsv2                  == other.rsv2 and
+         self.rsv3                  == other.rsv3 and
+         self.opcode                == other.opcode and
+         self.mask_bit              == other.mask_bit and
+         self.payload_length_code   == other.payload_length_code and
+         self.masking_key           == other.masking_key and
+         self.payload               == other.payload and
+         self.decoded_payload       == other.decoded_payload and
+         self.actual_payload_length == other.actual_payload_length)
+
 def apply_mask(message, masking_key):
     """
     Data sent from the server must be masked to prevent malicious clients
@@ -300,16 +313,14 @@ def create_client_handshake(host, port, key, version, resource):
     request = "GET %s HTTP/1.1" % resource
     return build_handshake(headers, request)
 
-
-def create_server_handshake(key, magic = websockets_magic):
+def create_server_handshake(key):
     """
       The server response is a valid HTTP 101 response. 
     """ 
-    digest = b64encode(sha1(key + magic).hexdigest().decode('hex'))
     headers = [
         ('Connection', 'Upgrade'),
         ('Upgrade', 'websocket'),
-        ('Sec-WebSocket-Accept', digest)
+        ('Sec-WebSocket-Accept', create_server_nounce(key))
     ]
     request = "HTTP/1.1 101 Switching Protocols"
     return build_handshake(headers, request)
@@ -322,7 +333,6 @@ def build_handshake(headers, request):
     handshake.append(b'\r\n')
     return b'\r\n'.join(handshake)
 
-
 def read_handshake(read_bytes, num_bytes_per_read):
     """
       From provided function that reads bytes, read in a 
@@ -355,13 +365,26 @@ def get_payload_length_pair(payload_bytestring):
         length_code = 127
     return (length_code, actual_length)
 
-def server_process_handshake(handshake):
-    headers = Message(StringIO(handshake.split('\r\n', 1)[1]))
+def process_handshake_from_client(handshake):
+    headers = headers_from_http_message(handshake)
     if headers.get("Upgrade", None) != "websocket":
         return
     key = headers['Sec-WebSocket-Key']
     return key
 
-def generate_client_nounce():
+def process_handshake_from_server(handshake, client_nounce):
+    headers = headers_from_http_message(handshake)
+    if headers.get("Upgrade", None) != "websocket":
+        return
+    key = headers['Sec-WebSocket-Accept']
+    return key
+
+def headers_from_http_message(http_message):
+    return Message(StringIO(http_message.split('\r\n', 1)[1]))
+
+def create_server_nounce(client_nounce):
+    return b64encode(sha1(client_nounce + websockets_magic).hexdigest().decode('hex'))
+
+def create_client_nounce():
     return b64encode(os.urandom(16)).decode('utf-8')
 
diff --git a/test/test_websockets.py b/test/test_websockets.py
index 0b2647ef3..a5ebf3d1b 100644
--- a/test/test_websockets.py
+++ b/test/test_websockets.py
@@ -1,29 +1,78 @@
+from netlib import tcp
 from netlib import test
 from netlib.websockets import implementations as impl
 from netlib.websockets import websockets as ws
 import os
+from nose.tools import raises
 
 class TestWebSockets(test.ServerTestBase):
     handler = impl.WebSocketsEchoHandler
 
+    def random_bytes(self, n = 100):
+        return os.urandom(n)
+
     def echo(self, msg):
         client = impl.WebSocketsClient(("127.0.0.1", self.port))
         client.connect()
         client.send_message(msg)
         response = client.read_next_message()
-        print "Assert response: " + response + " == msg: " + msg
         assert response == msg
 
     def test_simple_echo(self):
         self.echo("hello I'm the client")
 
     def test_frame_sizes(self):
-        small_string     =  os.urandom(100)   # length can fit in the the 7 bit payload length 
-        medium_string    =  os.urandom(50000) # 50kb, sligthly larger than can fit in a 7 bit int
-        large_string     =  os.urandom(150000) # 150kb, slightly larger than can fit in a 16 bit int
+        small_msg  = self.random_bytes(100)   # length can fit in the the 7 bit payload length 
+        medium_msg = self.random_bytes(50000) # 50kb, sligthly larger than can fit in a 7 bit int
+        large_msg  = self.random_bytes(150000) # 150kb, slightly larger than can fit in a 16 bit int
+
+        self.echo(small_msg)
+        self.echo(medium_msg)
+        self.echo(large_msg)
+
+    def test_default_builder(self):
+        """
+          default builder should always generate valid frames
+        """ 
+        msg = self.random_bytes()
+        client_frame = ws.WebSocketsFrame.default(msg, from_client = True)
+        assert client_frame.is_valid()
+
+        server_frame = ws.WebSocketsFrame.default(msg, from_client = False)
+        assert server_frame.is_valid()
+
+    def test_serialization_bijection(self):
+        for is_client in [True, False]:
+          for num_bytes in [100, 50000, 150000]:
+            frame = ws.WebSocketsFrame.default(self.random_bytes(num_bytes), is_client)
+            assert frame == ws.WebSocketsFrame.from_bytes(frame.to_bytes())
+
+        bytes = b'\x81\x11cba'
+        assert ws.WebSocketsFrame.from_bytes(bytes).to_bytes() == bytes
+
+
+class BadHandshakeHandler(impl.WebSocketsEchoHandler):
+    def handshake(self):
+        client_hs = ws.read_handshake(self.rfile.read, 1)
+        key       = ws.process_handshake_from_client(client_hs)
+        response  = ws.create_server_handshake("malformed_key")
+        self.wfile.write(response)
+        self.wfile.flush()
+        self.handshake_done = True
+
+class TestBadHandshake(test.ServerTestBase):
+    """
+      Ensure that the client disconnects if the server handshake is malformed
+    """ 
+    handler = BadHandshakeHandler
+
+    @raises(tcp.NetLibDisconnect)
+    def test(self):
+        client = impl.WebSocketsClient(("127.0.0.1", self.port))
+        client.connect()
+        client.send_message("hello")
+
+
 
-        self.echo(small_string)
-        self.echo(medium_string)
-        self.echo(large_string)
 
 

From 0ed2a290639833d772b89cf333577820e84f8204 Mon Sep 17 00:00:00 2001
From: Chandler Abraham <cabraham@twitter.com>
Date: Sat, 11 Apr 2015 17:28:52 -0700
Subject: [PATCH 4/5] whitespace

---
 test/test_websockets.py | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/test/test_websockets.py b/test/test_websockets.py
index a5ebf3d1b..0c23e355f 100644
--- a/test/test_websockets.py
+++ b/test/test_websockets.py
@@ -70,9 +70,4 @@ class TestBadHandshake(test.ServerTestBase):
     def test(self):
         client = impl.WebSocketsClient(("127.0.0.1", self.port))
         client.connect()
-        client.send_message("hello")
-
-
-
-
-
+        client.send_message("hello")
\ No newline at end of file

From 2d72a1b6b56f1643cd1d8be59eee55aa7ca2f17f Mon Sep 17 00:00:00 2001
From: Chandler Abraham <cabraham@twitter.com>
Date: Mon, 13 Apr 2015 13:36:09 -0700
Subject: [PATCH 5/5] 100% test coverage, though still need plenty more

---
 netlib/http.py                       | 14 --------------
 netlib/websockets/implementations.py | 10 ++--------
 netlib/websockets/websockets.py      |  9 ++++-----
 test/test_websockets.py              | 14 ++++++++++++--
 4 files changed, 18 insertions(+), 29 deletions(-)

diff --git a/netlib/http.py b/netlib/http.py
index 2c72621dd..264388636 100644
--- a/netlib/http.py
+++ b/netlib/http.py
@@ -29,20 +29,6 @@ def _is_valid_host(host):
         return None
     return True
 
-def is_successful_upgrade(request, response):
-    """
-        determines if a client and server successfully agreed to an HTTP protocol upgrade
-
-        https://developer.mozilla.org/en-US/docs/Web/HTTP/Protocol_upgrade_mechanism
-    """
-    http_switching_protocols_code = 101
-    
-    if request and response:
-        responseUpgrade = request.headers.get("Upgrade")
-        requestUpgrade  = response.headers.get("Upgrade")
-        if response.code == http_switching_protocols_code and responseUpgrade == requestUpgrade:
-            return requestUpgrade[0] if len(requestUpgrade) > 0 else None
-    return None
 
 def parse_url(url):
     """
diff --git a/netlib/websockets/implementations.py b/netlib/websockets/implementations.py
index 73a846905..1ded3b857 100644
--- a/netlib/websockets/implementations.py
+++ b/netlib/websockets/implementations.py
@@ -65,9 +65,6 @@ class WebSocketsClient(tcp.TCPClient):
         self.wfile.flush()
 
         server_handshake = ws.read_handshake(self.rfile.read, 1)
-        
-        if not server_handshake:
-            self.close()
 
         server_nounce = ws.process_handshake_from_server(server_handshake, self.client_nounce)
 
@@ -75,11 +72,8 @@ class WebSocketsClient(tcp.TCPClient):
             self.close()
 
     def read_next_message(self):
-        try:
-            return ws.WebSocketsFrame.from_byte_stream(self.rfile.read).payload
-        except IndexError:
-            self.close()
- 
+        return ws.WebSocketsFrame.from_byte_stream(self.rfile.read).payload
+
     def send_message(self, message):
         frame = ws.WebSocketsFrame.default(message, from_client = True)
         self.wfile.write(frame.safe_to_bytes())
diff --git a/netlib/websockets/websockets.py b/netlib/websockets/websockets.py
index cf9a68aa9..ea3db21d5 100644
--- a/netlib/websockets/websockets.py
+++ b/netlib/websockets/websockets.py
@@ -158,11 +158,10 @@ class WebSocketsFrame(object):
           ("actual_payload_length - " + str(self.actual_payload_length))])
 
     def safe_to_bytes(self):
-      try:
-          assert self.is_valid()
-          return self.to_bytes()
-      except:
-          raise WebSocketFrameValidationException()
+        if self.is_valid():
+            return self.to_bytes()
+        else:
+            raise WebSocketFrameValidationException()
 
     def to_bytes(self):
         """
diff --git a/test/test_websockets.py b/test/test_websockets.py
index 0c23e355f..951aa41ff 100644
--- a/test/test_websockets.py
+++ b/test/test_websockets.py
@@ -22,8 +22,8 @@ class TestWebSockets(test.ServerTestBase):
         self.echo("hello I'm the client")
 
     def test_frame_sizes(self):
-        small_msg  = self.random_bytes(100)   # length can fit in the the 7 bit payload length 
-        medium_msg = self.random_bytes(50000) # 50kb, sligthly larger than can fit in a 7 bit int
+        small_msg  = self.random_bytes(100)    # length can fit in the the 7 bit payload length 
+        medium_msg = self.random_bytes(50000)  # 50kb, sligthly larger than can fit in a 7 bit int
         large_msg  = self.random_bytes(150000) # 150kb, slightly larger than can fit in a 16 bit int
 
         self.echo(small_msg)
@@ -42,6 +42,10 @@ class TestWebSockets(test.ServerTestBase):
         assert server_frame.is_valid()
 
     def test_serialization_bijection(self):
+        """
+          Ensure that various frame types can be serialized/deserialized back and forth
+          between to_bytes() and from_bytes()
+        """ 
         for is_client in [True, False]:
           for num_bytes in [100, 50000, 150000]:
             frame = ws.WebSocketsFrame.default(self.random_bytes(num_bytes), is_client)
@@ -50,6 +54,12 @@ class TestWebSockets(test.ServerTestBase):
         bytes = b'\x81\x11cba'
         assert ws.WebSocketsFrame.from_bytes(bytes).to_bytes() == bytes
 
+    @raises(ws.WebSocketFrameValidationException)
+    def test_safe_to_bytes(self):
+        frame = ws.WebSocketsFrame.default(self.random_bytes(8))
+        frame.actual_payload_length = 1 #corrupt the frame
+        frame.safe_to_bytes()
+
 
 class BadHandshakeHandler(impl.WebSocketsEchoHandler):
     def handshake(self):