实现一个以websocket协议进行通信的服务端

# encoding = utf-8

import socket
import base64
import hashlib

# 获取请求头
def get_headers(data):
    header_dict = dict()
    str_data = str(data, encoding='utf-8')
    header, body = str_data.split('\r\n\r\n', 1)
    header_list = header.split('\r\n')
    for i in range(0, len(header_list)):
        if i == 0:
            if len(header_list[i].split(" ")) == 3:
                header_dict['method'], header_dict['url'], header_dict['protocol'] = header_list[i].split(' ')
        else:
            k, v = header_list[i].split(':', 1)
            header_dict[k] = v.strip()
    return header_dict


# 握手
def handshake(ws_socket, headers):
    response_tp1 = "HTTP/1.1 101 Switching Protocols\r\n" \
                   "Upgrade:websocket\r\n" \
                   "Connection: Upgrade\r\n" \
                   "Sec-WebSocket-Accept: %s\r\n" \
                   "WebSocket-Location: ws://%s%s\r\n\r\n"
    magic_string = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
    value = headers['Sec-WebSocket-Key'] + magic_string
    ac = base64.b64encode(hashlib.sha1(value.encode('utf-8')).digest())
    response_str = response_tp1 % (ac.decode('utf-8'), headers['Host'], headers['url'])
    ws_socket.send(bytes(response_str, encoding='utf-8'))


# 接收消息
def recv_msg(client_socket, recv_data):
    try:
        headers = get_headers(recv_data)
    except:
        payload_len = recv_data[1] & 127
        if payload_len == 126:
            extend_payload_len = recv_data[2:4]
            mask = recv_data[4:8]
            decoded = recv_data[8:]
        elif payload_len == 127:
            extend_payload_len = recv_data[2:10]
            mask = recv_data[10:14]
            decoded = recv_data[14:]
        else:
            extend_payload_len = None
            mask = recv_data[2:6]
            decoded = recv_data[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')
        print(">"*50)
        print(body)
    else:
        if headers.get('Connection'):
            handshake(client_socket, headers)


# 主函数
def main():
    tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    tcp_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    host = "127.0.0.1"
    port = 5000
    tcp_socket.bind((host, port))
    tcp_socket.listen(128)
    tcp_socket.setblocking(False)
    client_list = []
    while True:
        try:
            new_client_socket, client_addr = tcp_socket.accept()
        except Exception as e:
            pass
        else:
            new_client_socket.setblocking(False)
            client_list.append(new_client_socket)
            print(str(client_addr))
        for client_socket in client_list:
            try:
                recv_data = client_socket.recv(8096)

            except Exception as e:
                pass
            else:
                if recv_data:
                    recv_msg(client_socket, recv_data)
                else:
                    client_list.remove(client_socket)
                    client_socket.close()


if __name__ == '__main__':
    main()