gistfile1.py

import sys, time
import socket, threading
import hashlib, base64, struct
from libs.Event import *

class Connection(EventSource):

    HANDLER_SHAKE_RESPONSE = \
    "HTTP/1.1 101 Web Socket Protocol Handshake\r\n"\
    "Upgrade: websocket\r\n"\
    "Connection: Upgrade\r\n"\
    "Sec-WebSocket-Accept: %s\r\n"\
    "WebSocket-Origin: %s\r\n"\
    "WebSocket-Location: ws://%s/\r\n"\
    "WebSocket-Protocol:WebManagerSocket\r\n\r\n"

    HANDLER_SHAKE_KEY = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'

    OPCODE_CLOSE = 0x8

    __uuid = int(time.time())

    def __init__(self, conn, addr):
        EventSource.__init__(self)
        self.dispatchEvent(
            Event("data"),
            Event("error"),
            Event("close")
        )
        self.__conn       = conn
        self.client       = addr[0]
        self.port         = addr[1]
        self.__thread     = threading.Thread(target = self.__receiveThread)
        Connection.__uuid = Connection.__uuid + 1
        self.__id         = str(Connection.__uuid)
        self.__handShake()

    def receive(self):
        self.__thread.start()

    def connection(self):
        return self.__conn
    def id(self):
        return self.__id

    def send(self, data):
        stream   = data.encode("utf-8")
        response = bytearray([0x81])
        datalen  = len(stream)
        if (datalen <= 125):
            pass
        elif (datalen <= 65536):
            response.append(126)
        elif (datalen > 125):
            response.append(127)
        response.append(datalen)
        for byte in stream:
            response.append(byte)
        self.__conn.send(response)

    def __handShake(self):
        shake = self.__conn.recv(1024)
        addr  = self.client + ":" + str(self.port)
        if (not len(shake)): 
            self.__conn.close()
            raise HandShakeError("Handshake Failed, Invalid Connection " + addr)

        headers = {}
        lines   = shake.splitlines()
        for line in lines:
            line = line.decode("utf-8")
            parts = line.split(": ", 1)
            if (len(parts) == 2) :
                headers[parts[0]] = parts[1]

        headers['code'] = lines[-1]

        if "Connection" not in headers or  headers["Connection"] != "Upgrade"\
        or "Upgrade" not in headers or headers["Upgrade"] != "websocket":
            self.__conn.close()
            raise HandShakeError("Handshake Failed, this socket is not websocket " + addr)
        
        response = Connection.HANDLER_SHAKE_RESPONSE % (
            base64.b64encode(
                hashlib.sha1((headers["Sec-WebSocket-Key"] + Connection.HANDLER_SHAKE_KEY).encode("utf-8")).digest()
            ).decode("utf-8"),
            headers["Sec-WebSocket-Origin"] if "Sec-WebSocket-Origin" in headers else "null",
            headers["Host"]
        )
        self.__conn.send(response.encode("utf-8"))

    def __receiveThread(self):
        addr = self.client + ":" + str(self.port)
        while True:
            try:
                stream  = self.__conn.recv(1024)
                if (not stream or len(stream) < 6): break
                code    = stream[0]
                FIN     = (0xFF & code) >> 7
                RSV     = (0x70 & code) >> 4
                RSV1    = RSV >> 2
                RSV2    = 0x2 & RSV >> 1
                RSV3    = 0x1 & RSV
                OPCODE  = 0x0A & code
                code    = stream[1]
                MASK    = (0x80 & code) >> 7
                datalen = 0x7F & code
                data    = None
                assert(RSV1 == 0 and RSV2 == 0 and RSV3 == 0)
                assert(OPCODE < 0xF)
            except Exception as error:
                self.fireEvent("error", [error])
                break;

            if (OPCODE == Connection.OPCODE_CLOSE):
                break;

            if (datalen == 126):
                mask_key  = stream[4:8]
                datalen   = byte2int(stream[2:4])
                mask_data = stream[8:8+datalen]
            elif (datalen == 127):
                mask_key  = stream[10:14]
                datalen   = byte2int(stream[2:10])
                mask_data = stream[14:14+datalen]
            else:
                mask_key  = stream[2:6]
                mask_data = stream[6:6+datalen]

            if (datalen > 0):
                data = bytearray([mask_data[i] ^ mask_key[i%4] for i in range(len(mask_data))])

            self.fireEvent("data", [data.decode("utf-8"), stream])

        self.fireEvent("close", [])
        
class Error(Exception):
    def __init__(self, value):
        self.__value = value

    def __str__(self):
        return self.__value

class HandShakeError(Error):pass

def byte2int(byte):
    r = 0
    l = len(byte)
    for v in byte:
        l -= 1
        r += v << l*8
    return r