- workable with echo.websocket.org

2 files changed, 73 insertions, 26 deletions

diff --git a/test_websocket.py b/test_websocket.py
index c32b631..191d282 100644
--- a/test_websocket.py
+++ b/test_websocket.py
@@ -6,6 +6,9 @@ import websocket as ws
 
 TRACABLE=False
 
+def create_mask_key(n):
+    return "abcd"
+
 class StringSockMock:
     def __init__(self):
         self.set_data("")
@@ -137,43 +140,51 @@ class WebSocketTest(unittest.TestCase):
         self.assertRaises(ws.WebSocketException, sock._read_headers)
 
     def testSend(self):
+        # TODO: add longer frame data
         sock = ws.WebSocket()
+        sock.set_mask_key(create_mask_key)
         s = sock.io_sock = sock.sock = HeaderSockMock("data/header01.txt")
         sock.send("Hello")
-        #self.assertEquals(s.sent[0], "\x00Hello\xff")
+        self.assertEquals(s.sent[0], "\x81\x85abcd)\x07\x0f\x08\x0e")
+
         sock.send("こんにちは")
-        #self.assertEquals(s.sent[1], "\x00こんにちは\xff")
+        self.assertEquals(s.sent[1], "\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc")
+
         sock.send(u"こんにちは")
-        #self.assertEquals(s.sent[1], "\x00こんにちは\xff")
+        self.assertEquals(s.sent[1], "\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc")
 
     def testRecv(self):
+        # TODO: add longer frame data
         sock = ws.WebSocket()
         s = sock.io_sock = sock.sock = StringSockMock()
-        s.set_data("\x00こんにちは\xff")
+        s.set_data("\x81\x8fabcd\x82\xe3\xf0\x87\xe3\xf1\x80\xe5\xca\x81\xe2\xc5\x82\xe3\xcc")
         data = sock.recv()
         self.assertEquals(data, "こんにちは")
         
-        s.set_data("\x81\x05Hello")
+        s.set_data("\x81\x85abcd)\x07\x0f\x08\x0e")
         data = sock.recv()
         self.assertEquals(data, "Hello")
 
-        s.set_data("\x81\x81\x7f" + ("a"*255))
-        data = sock.recv()
-        self.assertEquals(len(data), 255)
-        self.assertEquals(data, "a" * 255)
-
     def testWebSocket(self):
         s  = ws.create_connection("ws://echo.websocket.org/") #ws://localhost:8080/echo")
         self.assertNotEquals(s, None)
         s.send("Hello, World")
         result = s.recv()
         self.assertEquals(result, "Hello, World")
+
         s.send("こにゃにゃちは、世界")
         result = s.recv()
         self.assertEquals(result, "こにゃにゃちは、世界")
         s.close()
 
-    def testSecureWebsocket(self):
+    def testPingPong(self):
+        s  = ws.create_connection("ws://echo.websocket.org/")
+        self.assertNotEquals(s, None)
+        s.ping("Hello")
+        s.pong("Hi")
+        s.close()
+        
+    def testSecureWebSocket(self):
         s  = ws.create_connection("wss://echo.websocket.org/")
         self.assertNotEquals(s, None)
         self.assert_(isinstance(s.io_sock, ws._SSLSocketWrapper))
diff --git a/websocket.py b/websocket.py
index 9a1d052..bb7d675 100644
--- a/websocket.py
+++ b/websocket.py
@@ -22,6 +22,7 @@ Copyright (C) 2010 Hiroki Ohtani(liris)
 
 import socket
 from urlparse import urlparse
+import os
 import struct
 import uuid
 import sha
@@ -175,7 +176,7 @@ class ABNF(object):
     LENGTH_63 = 1 << 63
 
     def __init__(self, fin = 0, rsv1 = 0, rsv2 = 0, rsv3 = 0,
-                 opcode = OPCODE_TEXT, mask = 0, data = ""):
+                 opcode = OPCODE_TEXT, mask = 1, data = ""):
         self.fin = fin
         self.rsv1 = rsv1
         self.rsv2 = rsv2
@@ -183,12 +184,13 @@ class ABNF(object):
         self.opcode = opcode
         self.mask = mask
         self.data = data
+        self.get_mask_key = os.urandom
 
     @staticmethod
     def create_frame(data, opcode):
         if opcode == ABNF.OPCODE_TEXT and isinstance(data, unicode):
             data = data.encode("utf-8")
-        return ABNF(1, 0, 0, 0, opcode, 0, data)
+        return ABNF(1, 0, 0, 0, opcode, 1, data)
 
     def format(self):
         if not is_bool(self.fin, self.rsv1, self.rsv2, self.rsv3):
@@ -213,8 +215,24 @@ class ABNF(object):
         
         if not self.mask:
             return frame_header + self.data
-        
-        raise NotImplementedError("masked format is not implemented")
+        else:
+            mask_key = self.get_mask_key(4)
+            return frame_header + self._get_masked(mask_key)
+
+    def _get_masked(self, mask_key):
+        s = ABNF.mask(mask_key, self.data)
+        return mask_key + "".join(s)
+
+    @staticmethod
+    def mask(mask_key, data):
+        _m = map(ord, mask_key)
+        _d = map(ord, data)
+        for i in range(len(_d)):
+            _d[i] ^= _m[i % 4]
+        s = map(chr, _d)
+        return "".join(s)
+
+
         
         
 
@@ -242,6 +260,10 @@ class WebSocket(object):
         """
         self.connected = False
         self.io_sock = self.sock = socket.socket()
+        self.get_mask_key = None
+        
+    def set_mask_key(self, func):
+        self.get_mask_key = func
 
     def settimeout(self, timeout):
         """
@@ -361,25 +383,38 @@ class WebSocket(object):
         
         return status, headers    
     
-    def send(self, payload, binary = False):
+    def send(self, payload, opcode = ABNF.OPCODE_TEXT, binary = False):
         """
         Send the data as string. payload must be utf-8 string or unicoce.
         """
-        frame = ABNF.create_frame(payload, ABNF.OPCODE_TEXT)
+        frame = ABNF.create_frame(payload, opcode)
+        if self.get_mask_key:
+            frame.get_mask_key = self.get_mask_key
         data = frame.format()
-        print repr(data)
         self.io_sock.send(data)
         if traceEnabled:
             logger.debug("send: " + repr(data))
 
+    def ping(self, payload):
+        self.send(payload, ABNF.OPCODE_PING)
+
+    def pong(self, payload):
+        self.send(payload, ABNF.OPCODE_PONG)
+
     def recv(self):
         """
         Receive utf-8 string data from the server.
         """
-        frame = self.read_frame()
-        return frame.data
-
-    def read_frame(self):
+        while True:
+            frame = self.recv_frame()
+            if frame.opcode in (ABNF.OPCODE_TEXT, ABNF.OPCODE_BINARY):
+                return frame.data
+            elif frame.opcode == ABNF.OPCODE_CLOSE:
+                return None
+            elif frame.opcode == ABNF.OPCODE_PING:
+                self.pong("Hi!")
+
+    def recv_frame(self):
         header_bytes = self._recv(2)
         b1 = ord(header_bytes[0])
         fin = b1 >> 7 & 1
@@ -398,11 +433,12 @@ class WebSocket(object):
         elif length == 0x7f:
             l = self._recv(8)
             length = struct.unpack("!Q", l)[0]
-        
-        data = self._recv(length)
 
         if mask:
-            raise NotImplementedError("masked data transfer is not implemented")
+            mask_key = self._recv(4)
+        data = self._recv(length)
+        if mask:
+            data = ABNF.mask(mask_key, data)
         
         frame = ABNF(fin, rsv1, rsv2, rsv3, opcode, mask, data)
         return frame
@@ -414,7 +450,7 @@ class WebSocket(object):
         """
         if self.connected:
             try:
-                self.io_sock.send("\xff\x00")
+                self.send("bye", ABNF.OPCODE_CLOSE)
                 timeout = self.sock.gettimeout()
                 self.sock.settimeout(1)
                 try: