# coding: utf-8
import json
import socket
import struct
import hashlib
import base64
import threading
import time

import logger

connection_list = {}  # 存放链接客户fd,元组
g_code_length = 0
g_header_length = 0  # websocket数据头部长度
PRINT_FLAG = False
return_time = "2020-09-09 00:00:00"
lock = threading.RLock()

"""
经测试发现IE 11浏览器在成功建立websocket连接后，会间隔30s发送空信息给服务器以验证是否处于连接状态，
因此服务区需要对收到的数据进行解码并判断其中载荷内容是否为空，如为空，应不进行广播
"""


# 计算web端提交的数据长度并返回
def get_data_length(message):
    global g_code_length
    global g_header_length
    try:
        g_code_length = message[1] & 127
        if g_code_length == 126:
            g_code_length = struct.unpack('>H', message[2:4])[0]
            g_header_length = 8
        elif g_code_length == 127:
            g_code_length = struct.unpack('>Q', message[2:10])[0]
            g_header_length = 14
        else:
            g_header_length = 6
        g_code_length = int(g_code_length)
    except Exception:
        g_code_length = 0
    return g_code_length


# 解析web端提交的bytes信息，返回str信息（可以解析中文信息）
def parse_data(message):
    """
    :param message: 需要解析的入参消息
    :type message: bytes
    :return: 返回str信息（可以解析中文信息） 失败为空
    :rtype: str
    """
    global g_code_length
    try:
        g_code_length = message[1] & 127
        if g_code_length == 126:
            g_code_length = struct.unpack('>H', message[2:4])[0]
            masks = message[4:8]
            data = message[8:]
        elif g_code_length == 127:
            g_code_length = struct.unpack('>Q', message[2:10])[0]
            masks = message[10:14]
            data = message[14:]
        else:
            masks = message[2:6]
            data = message[6:]
        en_bytes = b""
        cn_bytes = []
        for i, d in enumerate(data):
            nv = chr(d ^ masks[i % 4])
            nv_bytes = nv.encode()
            nv_len = len(nv_bytes)
            if nv_len == 1:
                en_bytes += nv_bytes
            else:
                en_bytes += b'%s'
                cn_bytes.append(ord(nv_bytes.decode()))
        if len(cn_bytes) > 2:
            cn_str = ""
            clen = len(cn_bytes)
            count = int(clen / 3)
            for x in range(count):
                i = x * 3
                b = bytes([cn_bytes[i], cn_bytes[i + 1], cn_bytes[i + 2]])
                # try:
                cn_str += b.decode()
                # except Exception as e:
                #     print(e)
            new = en_bytes.replace(b'%s%s%s', b'%s')
            new = new.decode()
            res = (new % tuple(list(cn_str)))
        else:
            res = en_bytes.decode()
    except Exception as e:
        # print("解析消息出错：", message, e)
        # if not res:
        res = ""
    return res


# 调用socket的send方法发送str信息给web端
def sendMessage(message):
    """
    :param message: 消息内容
    :type message: str
    :return: 无
    :rtype: None
    """
    global connection_list
    global lock
    lock.acquire()
    try:
        send_msg = b""  # 使用bytes格式,避免后面拼接的时候出现异常
        send_msg += b"\x81"
        back_str = []
        back_str.append('\x81')
        data_length = len(message.encode())  # 可能有中文内容传入，因此计算长度的时候需要转为bytes信息
        if PRINT_FLAG:
            print("INFO: send message is %s and len is %d" % (message, len(message.encode('utf-8'))))
        # 数据长度的三种情况
        if data_length <= 125:  # 当消息内容长度小于等于125时，数据帧的第二个字节0xxxxxxx 低7位直接标示消息内容的长度
            send_msg += str.encode(chr(data_length))
        elif data_length <= 65535:  # 当消息内容长度需要两个字节来表示时,此字节低7位取值为126,由后两个字节标示信息内容的长度
            send_msg += struct.pack('b', 126)
            send_msg += struct.pack('>h', data_length)
        elif data_length <= (2 ^ 64 - 1):  # 当消息内容长度需要把个字节来表示时,此字节低7位取值为127,由后8个字节标示信息内容的长度
            send_msg += struct.pack('b', 127)
            send_msg += struct.pack('>q', data_length)
        else:
            print(u'太长了')
        send_message = send_msg + message.encode('utf-8')
        del_list = []
        if connection_list:
            for connection in connection_list.values():
                if send_message is not None and len(send_message) > 0:
                    try:
                        connection.send(send_message)
                    except ConnectionAbortedError:
                        print("找不到连接：准备删除", connection)
                        del_list.append([k for k, v in connection_list.items() if v == connection][0])
        if del_list:
            for i in del_list:
                print("找不到连接：正在删除", i)
                del connection_list[i]
    finally:
        lock.release()


def send_recv_message(message):  # 直接回复
    global connection_list
    global lock
    lock.acquire()
    try:
        del_list = []
        if connection_list:
            for connection in connection_list.values():
                if message is not None and len(message) > 0:
                    try:
                        connection.send(message)
                    except ConnectionAbortedError:
                        print("找不到连接：准备删除", connection)
                        del_list.append([k for k, v in connection_list.items() if v == connection][0])
        if del_list:
            for i in del_list:
                print("找不到连接：正在删除", i)
                del connection_list[i]
    finally:
        lock.release()


# 删除连接,从集合中删除连接对象item
def delete_connection(item):
    global connection_list
    del connection_list['connection' + item]


# 定义WebSocket对象(基于线程对象)
class WebSocket(threading.Thread):
    def __init__(self, conn, index, name, remote, path=""):
        # 初始化线程
        threading.Thread.__init__(self)
        # 初始化数据,全部存储到自己的数据结构中self
        self.conn = conn
        self.index = index
        self.name = name
        self.remote = remote
        self.path = path
        self.GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
        self.buffer = ""
        self.buffer_utf8 = b""
        self.length_buffer = 0
        self.head_sign = False  # Socket是否握手的标志,初始化为false
    
    def generate_token(self, WebSocketKey):
        WebSocketKey = WebSocketKey + self.GUID
        Ser_WebSocketKey = hashlib.sha1(WebSocketKey.encode(encoding='utf-8')).digest()
        WebSocketToken = base64.b64encode(Ser_WebSocketKey)  # 返回的是一个bytes对象
        return WebSocketToken.decode('utf-8')

    # 运行线程
    def run(self):
        # Log输出,套接字index启动
        if PRINT_FLAG:
            print('Socket %s Start!' % self.index)
        global g_code_length
        global g_header_length
        global return_time
        while True:
            if self.head_sign is False:  # 如果没有进行握手
                try:
                    if PRINT_FLAG:
                        print('INFO: Socket %s Start head sign with %s!' % (self.index, self.remote))
                    self.buffer = self.conn.recv(1024).decode(
                        'utf-8')  # socket会话收到的只能是utf-8编码的信息，将接收到的bytes数据，通过utf-8编码方式解码为unicode编码进行处理
                    if PRINT_FLAG:
                        print("INFO: Socket %s self.buffer is {%s}" % (self.index, self.buffer))
                    if self.buffer.find('\r\n\r\n') != -1:
                        headers = {}
                        header, data = self.buffer.split('\r\n\r\n', 1)  # 按照这种标志分割一次,结果为：header data
                        # 对header进行分割后，取出后面的n-1个部分
                        for line in header.split("\r\n")[1:]:  # 再对header 和 data部分进行单独的解析
                            key, value = line.split(": ", 1)  # 逐行的解析Request Header信息(Key,Value)
                            headers[key] = value
                        try:
                            WebSocketKey = headers["Sec-WebSocket-Key"]
                        except KeyError:
                            print("Socket %s Handshaken Failed!" % self.index)
                            delete_connection(str(self.index))
                            self.conn.close()
                            break
                        WebSocketToken = self.generate_token(WebSocketKey)
                        headers["Location"] = ("ws://%s%s" % (headers["Host"], self.path))
                        # 握手过程,服务器构建握手的信息,进行验证和匹配
                        # Upgrade: WebSocket 表示为一个特殊的http请求,请求目的为从http协议升级到websocket协议
                        handshake = "HTTP/1.1 101 Switching Protocols\r\n" \
                                    "Connection: Upgrade\r\n" \
                                    "Sec-WebSocket-Accept: " + WebSocketToken + "\r\n" \
                                                                                "Upgrade: websocket\r\n\r\n"
                        self.conn.send(handshake.encode(encoding='utf-8'))  # 前文以bytes类型接收，此处以bytes类型进行发送
                        # 此处需要增加代码判断是否成功建立连接
                        self.head_sign = True  # socket连接成功建立之后修改握手标志
                        # 向全部连接客户端集合发送消息,(环境套接字x的到来)
                        sendMessage("Welcome " + self.name + " !")
                        g_code_length = 0

                    else:
                        print("Socket %s Error2!" % self.index)
                        delete_connection(str(self.index))
                        self.conn.close()
                        break
                except Exception as e:
                    print(e, "握手连接失败正在删除 connection" + str(self.index))
                    delete_connection(str(self.index))
                    self.conn.close()
                    break
            else:
                # 每次接收1024字节数据，需要判断是否接收完所有数据，如没有接收完，需要循环接收完再处理
                try:
                    mm = self.conn.recv(1024)
                except Exception as e:
                    mm = b''
                    if '[WinError 10054] 远程主机强迫关闭了一个现有的连接。' in str(e):
                        delete_connection(str(self.index))
                        self.conn.close()
                        break
                # 计算接受的长度，判断是否接收完，如未接受完需要继续接收
                if mm != "b''":
                    if g_code_length == 0:
                        get_data_length(mm)  # 调用此函数可以计算并修改全局变量g_code_length和g_header_length的值
                    self.length_buffer += len(mm)
                    self.buffer_utf8 += mm
                    if self.length_buffer - g_header_length < g_code_length:
                        if PRINT_FLAG:
                            print("INFO: 数据未接收完,接续接受")
                        continue
                    else:
                        if PRINT_FLAG:
                            print("g_code_length:", g_code_length)
                            print("INFO Line 204: Recv信息 %s,长度为 %d:" % (self.buffer_utf8, len(self.buffer_utf8)))
                        if not self.buffer_utf8:
                            continue
                        recv_message = parse_data(self.buffer_utf8)  # 接受信息转换
                        # recv_message = self.buffer_utf8
                        nowTime = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
                        if not recv_message:
                            new_time = time.strptime(nowTime, '%Y-%m-%d %H:%M:%S')
                            old_time = time.strptime(return_time, '%Y-%m-%d %H:%M:%S')
                            loop_time = time.gmtime(time.mktime(new_time) - time.mktime(old_time))
                            send_recv_message(self.buffer_utf8)
                            if int(loop_time.tm_sec) >= 30:
                                print("收到心跳并已回复" + nowTime)
                                return_time = nowTime
                        elif recv_message:
                            if recv_message == "quit":
                                print("Socket %s Logout!" % self.index)
                                nowTime = time.strftime('%H:%M:%S', time.localtime(time.time()))
                                sendMessage("%s %s say: %s" % (nowTime, self.remote, self.name + " Logout"))
                                delete_connection(str(self.index))
                                self.conn.close()
                                break
                            else:
                                try:
                                    json_str = json.loads(recv_message)
                                    if json_str['EventType'] != 0 and json_str['EventType'] != 1:
                                        if json_str['EventType'] == 1005 or json_str['EventType'] == 3006:
                                            recv_log = logger.logs("recv")
                                            recv_log.info(str(recv_message))
                                        else:
                                            print(nowTime, recv_message)
                                    else:
                                        if return_time:
                                            new_time = time.strptime(nowTime, '%Y-%m-%d %H:%M:%S')
                                            old_time = time.strptime(return_time, '%Y-%m-%d %H:%M:%S')
                                            loop_time = time.gmtime(time.mktime(new_time) - time.mktime(old_time))
                                            sendMessage(json_str)
                                            if int(loop_time.tm_sec) >= 6:
                                                if json_str['EventType'] == 1:
                                                    print("收到并已回复心跳：", nowTime, recv_message)
                                                    return_time = nowTime
                                            if json_str['InfoCode'] is not None:
                                                print(nowTime, recv_message)
                                except Exception as e:
                                    if e == '[WinError 10054] 远程主机强迫关闭了一个现有的连接。':
                                        delete_connection(str(self.index))
                                        self.conn.close()
                                        break
                                    # print("解析错误原因是：", e, "错误的消息是：", recv_message)
                                    new_time = time.strptime(nowTime, '%Y-%m-%d %H:%M:%S')
                                    old_time = time.strptime(return_time, '%Y-%m-%d %H:%M:%S')
                                    loop_time = time.gmtime(time.mktime(new_time) - time.mktime(old_time))
                                    # send_recv_message(recv_message)
                                    # send_recv_message(recv_message)
                                    send_recv_message(self.buffer_utf8)
                                    if int(loop_time.tm_sec) >= 30:
                                        print("收到心跳并已回复" + nowTime)
                                        return_time = nowTime
                        print(recv_message)
                        g_code_length = 0
                        self.length_buffer = 0
                        self.buffer_utf8 = b""


# WebSocket服务器对象()
class WebSocketServer(object):
    def __init__(self):
        self.socket = None
        self.i = 0
    
    # 开启操作
    def begin(self):
        if PRINT_FLAG:
            print('WebSocketServer Start!')
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        ip = '192.168.11.103'
        port = 8081
        if PRINT_FLAG:
            print("WebServer is listening %s,%d" % (ip, port))
        self.socket.bind((ip, port))
        self.socket.listen(50)
        # 全局连接集合
        global connection_list
        
        while True:
            # 服务器响应请求,返回连接客户的信息(连接fd,客户地址)
            connection, address = self.socket.accept()
            # 根据连接的客户信息,创建WebSocket对象(本质为一个线程)
            # sockfd，index，用户名，地址
            newSocket = WebSocket(connection, self.i, address[0], address)
            # 线程启动
            newSocket.start()
            # 更新连接的集合(hash表的对应关系)-name->sockfd
            connection_list['connection' + str(self.i)] = connection
            print("新连接", connection)
            self.i += 1


if __name__ == "__main__":
    server = WebSocketServer()
    # server.begin()
    t = threading.Thread(target=server.begin, name="线程消息", args=())
    print("ws启动了")
    t.start()
    while True:
        # if ws.state == -1:
        #     break
        msg = str(input("请输入需要发送的信息："))
        # ws.send_all(bytes(json.dumps(msg), encoding="utf8"))
        if msg:
            sendMessage(msg)
            send_log = logger.logs("send")
            send_log.info(str(msg))
        # print("服务端已发送的消息", msg)
        # time.sleep(5)
