From 1b509d5aea31a636b6c8ce854e0dd685e34d03de Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Tue, 21 Apr 2015 22:51:01 +1200
Subject: [PATCH] Whitespace, interface simplification

- safe_tobytes doesn't buy us much
- move masking key generation inline
---
 netlib/websockets.py    | 17 ++---------------
 test/test_websockets.py | 13 +++----------
 2 files changed, 5 insertions(+), 25 deletions(-)

diff --git a/netlib/websockets.py b/netlib/websockets.py
index a03185fae..0cd4dba15 100644
--- a/netlib/websockets.py
+++ b/netlib/websockets.py
@@ -29,10 +29,6 @@ class CONST(object):
     MAX_64_BIT_INT = (1 << 64)
 
 
-class WebSocketFrameValidationException(Exception):
-    pass
-
-
 class Frame(object):
     """
         Represents one websockets frame.
@@ -95,7 +91,8 @@ class Frame(object):
 
         if from_client:
             mask_bit = 1
-            masking_key = random_masking_key()
+            # Random masking key
+            masking_key = os.urandom(4)
             payload = apply_mask(message, masking_key)
         else:
             mask_bit = 0
@@ -164,12 +161,6 @@ class Frame(object):
         """
         return cls.from_file(io.BytesIO(bytestring))
 
-    def safe_to_bytes(self):
-        if self.is_valid():
-            return self.to_bytes()
-        else:
-            raise WebSocketFrameValidationException()
-
     def to_bytes(self):
         """
             Serialize the frame back into the wire format, returns a bytestring
@@ -308,10 +299,6 @@ def apply_mask(message, masking_key):
     return result
 
 
-def random_masking_key():
-    return os.urandom(4)
-
-
 def client_handshake_headers(key=None, version=VERSION):
     """
         Create the headers for a valid HTTP upgrade request. If Key is not
diff --git a/test/test_websockets.py b/test/test_websockets.py
index 9b27e810d..3fc67dfee 100644
--- a/test/test_websockets.py
+++ b/test/test_websockets.py
@@ -1,5 +1,4 @@
-from netlib import tcp, test, websockets, http, odict
-import io
+from netlib import tcp, test, websockets, http
 import os
 from nose.tools import raises
 
@@ -120,17 +119,11 @@ class TestWebSockets(test.ServerTestBase):
                     self.random_bytes(num_bytes), is_client
                 )
                 assert frame == websockets.Frame.from_bytes(
-                    frame.safe_to_bytes()
+                    frame.to_bytes()
                 )
 
         bytes = b'\x81\x03cba'
-        assert websockets.Frame.from_bytes(bytes).safe_to_bytes() == bytes
-
-    @raises(websockets.WebSocketFrameValidationException)
-    def test_safe_to_bytes(self):
-        frame = websockets.Frame.default(self.random_bytes(8))
-        frame.actual_payload_length = 1 # corrupt the frame
-        frame.safe_to_bytes()
+        assert websockets.Frame.from_bytes(bytes).to_bytes() == bytes
 
     def test_check_server_handshake(self):
         resp = http.Response(