import socket 
import hashlib
import socket
import base64
from threading import Thread
import time
#The websocket thread
class WebSocketStart(Thread):
    def __init__(self,connect):
        Thread.__init__(self)
        self.connect = connect
        self.isHandleShake = False
    # the thread method
    def run(self):
        while True:
            if not self.isHandleShake:
                # 第一次连接，需要握手
                header=self.getHeader()
                secKey=header['Sec-WebSocket-Key']
                acceptKey=self.generateAcceptKey(secKey)
                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"%(acceptKey.decode('utf-8'))
                self.connect.send(response.encode())
                self.isHandleShake=True
                print("HandleShake with client is Finish.")
            else:
                command = self.receiveMsg()
                print(command)
                self.sendMsgToClient("I have receive you msg")
    # get the handleshake's header
    def getHeader(self):
        requestData = self.connect.recv(1024).decode()
        requestList = requestData.split('\r\n')
        headers={}
        for reqItem in requestList:
            if ': ' in reqItem:
                unit = reqItem.split(': ')
                headers[unit[0]]=unit[1]
        return  headers
    # get the AcceptKey
    def generateAcceptKey(self,secKey):
        sha1 = hashlib.sha1()
        sha1.update((secKey+'258EAFA5-E914-47DA-95CA-C5AB0DC85B11').encode())
        sha1_result = sha1.digest()
        acceptKey = base64.b64encode(sha1_result)
        return  acceptKey
    # send message to client
    def sendMsgToClient(self,msg):
        import struct
        sendMsg = struct.pack('!B',0x81)
        length=len(msg)
        if length <= 125:
            sendMsg += struct.pack('!B', length)
        elif length <= 65536:
            sendMsg += struct.pack('!B', 126)
            sendMsg += struct.pack('!H', length)
        elif length == 127:
            sendMsg += struct.pack('!B', 127)
            sendMsg += struct.pack('!Q', length)
        sendMsg += struct.pack('!%ds' % (length), msg.encode())
        self.connect.send(sendMsg)
    # receive msy from client
    def receiveMsg(self):
        info=self.connect.recv(8096)
        payload_len = info[1] & 127
        if payload_len == 126:
            extend_payload_len = info[2:4]
            mask = info[4:8]
            decoded = info[8:]
        elif payload_len == 127:
            extend_payload_len = info[2:10]
            mask = info[10:14]
            decoded = info[14:]
        else:
            extend_payload_len = None
            mask = info[2:6]
            decoded = info[6:]
        bytes_list = bytearray()
        for i in range(len(decoded)):
            chunk = decoded[i] ^ mask[i % 4]
            bytes_list.append(chunk)
        body = str(bytes_list, encoding='utf-8')
        return body

def main():
    socke = socket.socket()
    socke.bind(('127.0.0.1',8908))
    socke.listen(5)
    while True:
        try:
            connect,address = socke.accept()
            WebSocketStart(connect).start()
        except:
            time.sleep(1)
if __name__ == '__main__':
    main()