jkp · July 18, 2012 13:28 · Jul 19, 2012 · Jul 18, 2012
diff --git a/websocketserver.py b/websocketserver.py
@@ -5,10 +5,9 @@
 from mimetools import Message
 from StringIO import StringIO
 
-PORT = 9999
-MAGIC = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
-
 class WebSocketsHandler(SocketServer.StreamRequestHandler):
+    magic = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
+
     def setup(self):
         SocketServer.StreamRequestHandler.setup(self)
         print "connection established", self.client_address
@@ -31,7 +30,20 @@ def read_next_message(self):
         decoded = ""
         for char in self.rfile.read(length):
             decoded += chr(ord(char) ^ masks[len(decoded) % 4])
-        self.process_message(decoded)
+        self.on_message(decoded)
+
+    def send_message(self, message):
+        self.request.send(chr(129))
+        length = len(message)
+        if length <= 125:
+            self.request.send(chr(length))
+        elif length >= 126 and length <= 65535:
+            self.request.send(126)
+            self.request.send(struct.pack(">H", length))
+        else:
+            self.request.send(127)
+            self.request.send(struct.pack(">Q", length))
+        self.request.send(message)
 
     def handshake(self):
         data = self.request.recv(1024).strip()
@@ -40,17 +52,17 @@ def handshake(self):
             return
         print 'Handshaking...'
         key = headers['Sec-WebSocket-Key']
-        digest = b64encode(sha1(key + MAGIC).hexdigest().decode('hex'))
+        digest = b64encode(sha1(key + self.magic).hexdigest().decode('hex'))
         response = 'HTTP/1.1 101 Switching Protocols\r\n'
         response += 'Upgrade: websocket\r\n'
         response += 'Connection: Upgrade\r\n'
         response += 'Sec-WebSocket-Accept: %s\r\n\r\n' % digest
         self.handshake_done = self.request.send(response)
 
-    def process_message(self, message):
+    def on_message(self, message):
         print messsage
 
 if __name__ == "__main__":
     server = SocketServer.TCPServer(
-        ("localhost", PORT), WebSocketsHandler)
-    server.serve_forever()
+        ("localhost", 9999), WebSocketsHandler)
+    server.serve_forever()
diff --git a/websocketserver.py b/websocketserver.py
@@ -0,0 +1,56 @@
+import struct
+import SocketServer
+from base64 import b64encode
+from hashlib import sha1
+from mimetools import Message
+from StringIO import StringIO
+
+PORT = 9999
+MAGIC = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
+
+class WebSocketsHandler(SocketServer.StreamRequestHandler):
+    def setup(self):
+        SocketServer.StreamRequestHandler.setup(self)
+        print "connection established", self.client_address
+        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):
+        length = ord(self.rfile.read(2)[1]) & 127
+        if length == 126:
+            length = struct.unpack(">H", self.rfile.read(2))[0]
+        elif length == 127:
+            length = struct.unpack(">Q", self.rfile.read(8))[0]
+        masks = [ord(byte) for byte in self.rfile.read(4)]
+        decoded = ""
+        for char in self.rfile.read(length):
+            decoded += chr(ord(char) ^ masks[len(decoded) % 4])
+        self.process_message(decoded)
+
+    def handshake(self):
+        data = self.request.recv(1024).strip()
+        headers = Message(StringIO(data.split('\r\n', 1)[1]))
+        if headers.get("Upgrade", None) != "websocket":
+            return
+        print 'Handshaking...'
+        key = headers['Sec-WebSocket-Key']
+        digest = b64encode(sha1(key + MAGIC).hexdigest().decode('hex'))
+        response = 'HTTP/1.1 101 Switching Protocols\r\n'
+        response += 'Upgrade: websocket\r\n'
+        response += 'Connection: Upgrade\r\n'
+        response += 'Sec-WebSocket-Accept: %s\r\n\r\n' % digest
+        self.handshake_done = self.request.send(response)
+
+    def process_message(self, message):
+        print messsage
+
+if __name__ == "__main__":
+    server = SocketServer.TCPServer(
+        ("localhost", PORT), WebSocketsHandler)
+    server.serve_forever()