#!/usr/bin/env python

import asyncore
import socket
import struct
import time
from hashlib import sha1
from base64 import encodestring


class WebSocketConnection(asyncore.dispatcher_with_send):

    TEXT = 0x01
    BINARY = 0x02

    def __init__(self, conn, server):
        asyncore.dispatcher_with_send.__init__(self, conn)

        self.server = server
        self.server.sessions.append(self)
        self.readystate = "connecting"
        self.buffer = ""

    def handle_read(self):
        data = self.recv(1024)
        self.buffer += data
        if self.readystate == "connecting":
            self.parse_connecting()
        elif self.readystate == "open":
            self.parse_frame()

    def handle_close(self):
        self.server.sessions.remove(self)
        self.close()

    def parse_connecting(self):
        """
        Parsowanie uscisku dloni WebSocket. Nie jest to pelny parser HTTP!
        """
        header_end = self.buffer.find("\r\n\r\n")
        if header_end == -1:
            return
        else:
            header = self.buffer[:header_end]
            # usuniecie naglowka i czterech koncowych bajtow z bufora
            self.buffer = self.buffer[header_end + 4:]
            header_lines = header.split("\r\n")
            headers = {}

            # walidacja zadania HTTP i utworzenie lokacji
            method, path, protocol = header_lines[0].split(" ")
            if method != "GET" or protocol != "HTTP/1.1" or path[0] != "/":
                self.terminate()
                return

            # parsowanie naglowka
            for line in header_lines[1:]:
                key, value = line.split(": ")
                headers[key] = value

            headers["Location"] = "ws://" + headers["Host"] + path

            self.readystate = "open"
            self.handler = self.server.handlers.get(path, None)(self)

            self.send_server_handshake_10(headers)

    def terminate(self):
        self.ready_state = "closed"
        self.close()

    def send_server_handshake_10(self, headers):
        """
        Przeslanie odpowiedzi na uscisk dloni zgodnej ze szkicem HyBi-10
        """
        key = headers["Sec-WebSocket-Key"]

        # wyswietlenie naglowka odpowiedzi
        self.send_bytes("HTTP/1.1 101 Switching Protocols\r\n")
        self.send_bytes("Upgrade: WebSocket\r\n")
        self.send_bytes("Connection: Upgrade\r\n")
        self.send_bytes("Sec-WebSocket-Accept: %s\r\n" % self.hash_key(key))

        if "Sec-WebSocket-Protocol" in headers:
            protocol = headers["Sec-WebSocket-Protocol"]
            self.send_bytes("Sec-WebSocket-Protocol: %s\r\n" % protocol)

    def hash_key(self, key):
        guid = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
        combined = key + guid
        hashed = sha1(combined).digest()
        return encodestring(hashed)

    def parse_frame(self):
        """
        Parsowanie ramki WebSocket. Jezeli w buforze nie ma kompletnej ramki,
		zakoncz bez modyfikowania bufora.
        """
        buf = self.buffer
        payload_start = 2

        # sprobuj pobrac pierwsze dwa bajty
        if len(buf) < 3:
            return
        b = ord(buf[0])
        fin = b & 0x80      # pierwszy bit
        # kolejne 3 bity zarezerwowane
        opcode = b & 0x0f   # niskie 4 bity
        b2 = ord(buf[1])
        mask = b2 & 0x80    # wysoki bit drugiego bajta
        length = b2 & 0x7f    # niskie 7 bitow drugiego bajta

        # sprawdz, czy pozostalo wystarczajaco duzo bajtow
        if len(buf) < payload_start + 4:
            return
        elif length == 126:
            length, = struct.unpack(">H", buf[2:4])
            payload_start += 2
        elif length == 127:
            length, = struct.unpack(">I", buf[2:6])
            payload_start += 4

        if mask:
            mask_bytes = [ord(b) for b in buf[payload_start:payload_start + 4]]
            payload_start += 4

        # czy ramka w buforze jest kompletna?
        if len(buf) < payload_start + length:
            return

        # usuniecie poczatkowych bajtow, dekodowanie, jezeli jest konieczne, przeslanie
        payload = buf[payload_start:payload_start + length]
        self.buffer = buf[payload_start + length:]

        # wykorzystanie xor i maskowanie bajtow w celu odmaskowania danych
        if mask:
            unmasked = [mask_bytes[i % 4] ^ ord(b)
                            for b, i in zip(payload, range(len(payload)))]
            payload = "".join([chr(c) for c in unmasked])

        if opcode == WebSocketConnection.TEXT:
            s = payload.decode("UTF8")
            self.handler.dispatch(s)
        if opcode == WebSocketConnection.BINARY:
            self.handler.dispatch(payload)
        return True

    def send(self, s):
        """
        Zakodowanie i przeslanie wiadomosci WebSocket
        """

        message = ""
        # zawsze przeslij cala wiadomosc w jednej ramce (fin)
        b1 = 0x80

        # w Pythonie 2 str sa bajtami, a kody unicode sa ciagami znakow
        if type(s) == unicode:
            b1 |= WebSocketConnection.TEXT
            payload = s.encode("UTF8")
        elif type(s) == str:
            b1 |= WebSocketConnection.BINARY
            payload = s

        message += chr(b1)

        # nigdy nie maskuj ramek przesylanych z serwera do klienta
        b2 = 0
        length = len(payload)
        if length < 126:
            b2 |= length
            message += chr(b2)
        elif length < (2 ** 16) - 1:
            b2 |= 126
            message += chr(b2)
            l = struct.pack(">H", length)
            message += l
        else:
            l = struct.pack(">Q", length)
            b2 |= 127
            message += chr(b2)
            message += l

        message += payload

        if self.readystate == "open":
            self.send_bytes(message)

    def send_bytes(self, bytes):
        try:
            asyncore.dispatcher_with_send.send(self, bytes)
        except:
            pass


class EchoHandler(object):
    """
    Klasa EchoHandler powtarza kazdy przychodzacy ciag znakow do tego samego polaczenia WebSocket.
    """

    def __init__(self, conn):
        self.conn = conn

    def dispatch(self, data):
        try:
            self.conn.send(data)
        except:
            pass


class WebSocketServer(asyncore.dispatcher):

    def __init__(self, port=80, handlers=None):
        asyncore.dispatcher.__init__(self)
        self.handlers = handlers
        self.sessions = []
        self.port = port
        self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
        self.set_reuse_addr()
        self.bind(("", port))
        self.listen(5)

    def handle_accept(self):
        conn, addr = self.accept()
        session = WebSocketConnection(conn, self)

if __name__ == "__main__":
    print "Start serwera WebSocket"
    WebSocketServer(port=8080, handlers={"/echo": EchoHandler})
    asyncore.loop()
