websockets: refactor implementation and add tests

1 files changed, 33 insertions, 0 deletions

diff --git a/netlib/websockets/masker.py b/netlib/websockets/masker.py
new file mode 100644
index 00000000..bd39ed6a
--- /dev/null
+++ b/netlib/websockets/masker.py
@@ -0,0 +1,33 @@
+from __future__ import absolute_import
+
+import six
+
+
+class Masker(object):
+    """
+    Data sent from the server must be masked to prevent malicious clients
+    from sending data over the wire in predictable patterns.
+
+    Servers do not have to mask data they send to the client.
+    https://tools.ietf.org/html/rfc6455#section-5.3
+    """
+
+    def __init__(self, key):
+        self.key = key
+        self.offset = 0
+
+    def mask(self, offset, data):
+        result = bytearray(data)
+        for i in range(len(data)):
+            if six.PY2:
+                result[i] ^= ord(self.key[offset % 4])
+            else:
+                result[i] ^= self.key[offset % 4]
+            offset += 1
+        result = bytes(result)
+        return result
+
+    def __call__(self, data):
+        ret = self.mask(self.offset, data)
+        self.offset += len(ret)
+        return ret