Python-websocket及心跳检测

1、安装websocket-client

pip install websocket-client

2、websocket服务端

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
@author: wdj
@contact:wei7wei@126.com
@file: ws.py
@time: 2022/6/8 9:34
"""
# coding=utf-8
import socket
import time
import hashlib
import base64
import struct
from multiprocessing import Process
HTTP_RESPONSE = "HTTP/1.1 {code} {msg}\r\n"  \
                "Server:LyricTool\r\n" \
                "Date:{date}\r\n" \
                "Content-Length:{length}\r\n" \
                "\r\n" \
                "{content}\r\n"
STATUS_CODE = {200: 'OK', 501: 'Not Implemented'}
UPGRADE_WS = "HTTP/1.1 101 Switching Protocols\r\n" \
             "Connection: Upgrade\r\n" \
             "Upgrade: websocket\r\n" \
             "Sec-WebSocket-Accept: {}\r\n" \
             "WebSocket-Protocol: chat\r\n\r\n"
def sec_key_gen(msg):
    key = msg + '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
    ser_key = hashlib.sha1(key.encode('utf-8')).digest()
    return base64.b64encode(ser_key).decode()
class WebsocketServer:
    def __init__(self, conn):
        # 接受一个socket对象
        self.conn = conn
        self.state = 0
    def open(self):
        self._handshake()
        if self.state == 1:
            return self
        else:
            raise Exception('Handsake failed.')
    def __enter__(self):
        return self.open()
    def getstate(self):
        # 获取连接状态
        state_map = {0: 'READY', 1: 'CONNECTION ESTABLISHED', 2: 'HANDSHAKED', 3: 'FAILED', -1: 'CLOSED'}
        return self.state, state_map[self.state]
    def _handshake(self):
        raw_data = b''
        while True:
            fragment = self.conn.recv(1024)
            raw_data += fragment
            if len(fragment) < 1024:
                break
        data = raw_data.decode('utf-8')
        header, content = data.split('\r\n\r\n', 1)
        header = header.split('\r\n')
        options = map(lambda i: i.split(': '), header[1:])
        options_dict = {item[0]: item[1] for item in options}
        date = time.strftime("%m,%d%Y", time.localtime())
        if 'Sec-WebSocket-Key' not in options_dict:
            self.conn.send(
                bytes(HTTP_RESPONSE.format(code=501, msg=STATUS_CODE[501], date=date, length=len(date), content=date),
                      encoding='utf-8'))
            self.conn.close()
            self.state = 3
            return True
        else:
            self.state = 2
            self._build(options_dict['Sec-WebSocket-Key'])
            return True
    def _build(self, sec_key):
        # 建立WebSocket连接
        response = UPGRADE_WS.format(sec_key_gen(sec_key))
        self.conn.send(bytes(response, encoding='utf-8'))
        self.state = 1
        return True
    def _get_data(self, info, setcode):
        payload_len = info[1] & 127
        fin = 1 if info[0] & 128 == 128 else 0
        opcode = info[0] & 15  # 提取opcode
        # 提取载荷数据
        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)
        if opcode == 0x00:
            opcode = setcode
        if opcode == 0x01:   # 文本帧
            body = str(bytes_list, encoding='utf-8')
            return fin, opcode, body
        elif opcode == 0x08:
            self.close()
            raise IOError('Connection closed by Client.')
        else:  # 二进制帧或其他，原样返回
            body = decoded
            return fin, opcode, body
    def recv(self):
        msg = self.conn.recv()
        print(msg)
        # 处理切片
        opcode = 0x00
        # while True:
        #     raw_data = b''
            # while True:
            #     section = self.conn.recv(1024)
            #     raw_data += section
            #     if len(section) < 1024:
            #         break
            # fin, _opcode, fragment = self._get_data(raw_data, opcode)
            # opcode = _opcode if _opcode != 0x00 else opcode
            # msg += fragment
            # if fin == 1:   # 是否是最后一个分片
            #     break
        return msg
    def send(self, msg, fin=True):
        # 发送数据
        data = struct.pack('B', 129) if fin else struct.pack('B', 0)
        msg_len = len(msg)
        if msg_len <= 125:
            data += struct.pack('B', msg_len)
        elif msg_len <= (2**16 - 1):
            data += struct.pack('!BH', 126, msg_len)
        elif msg_len <= (2**64 - 1):
            data += struct.pack('!BQ', 127, msg_len)
        else:
            # 分片传输超大内容（应该用不到）
            while True:
                fragment = msg[:(2**64 - 1)]
                msg -= fragment
                if msg > (2**64 - 1):
                    self.send(fragment, False)
                else:
                    self.send(fragment)
        data += bytes(msg, encoding='utf-8')
        self.conn.send(data)
    def ping(self):
        ping_msg = 0b10001001
        data = struct.pack('B', ping_msg)
        data += struct.pack('B', 0)
        while True:
            self.conn.send(data)
            data = self.conn.recv(1024)
            pong = data[0] & 127
            if pong != 9:
                self.close()
                raise IOError('Connection closed by Client.')
    def close(self):
        self.conn.close()
        self.state = -1
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is IOError:
            print(exc_val)
        self.close()
def ws_handler(conn):
    print("ws_handler")
    with WebsocketServer(conn) as ws:
        while True:
            time.sleep(1)
            # msg = ws.recv()
            # if ws.state == -1:
            #     break
            # print(msg)
            ws.send("999999")
            print("發送成功")
if __name__ == '__main__':
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind(('127.0.0.1',8087))
    s.listen(1)
    print('Server Started.')
    while True:
        con, addr = s.accept()
        print("Accepted. {0}, {1}".format(con, str(addr)))
        p = Process(target=ws_handler, args=(con,))
        p.start()

3、websocket客户端

import json
import websocket    # pip install websocket-client

CHANNELS_WS = [
    # 这里输入需要订阅的频道
]


class Feed(object):

    def __init__(self):
        self.url = 'ws://127.0.0.1:8087'      # 这里输入websocket的url
        self.ws = None

    def on_open(self, ws):
        """
        Callback object which is called at opening websocket.
        1 argument:
        @ ws: the WebSocketApp object
        """
        print('A new WebSocketApp is opened!')

        # 开始订阅（举个例子）
        sub_param = {"op": "subscribe", "args": CHANNELS_WS}
        sub_str = json.dumps(sub_param)
        ws.send(sub_str)
        print("Following Channels are subscribed!")
        print(CHANNELS_WS)

    def on_data(self, ws, string, type, continue_flag):
        """
        4 argument.
        The 1st argument is this class object.
        The 2nd argument is utf-8 string which we get from the server.
        The 3rd argument is data type. ABNF.OPCODE_TEXT or ABNF.OPCODE_BINARY will be came.
        The 4th argument is continue flag. If 0, the data continue
        """

    def on_message(self, ws, message):
        """
        Callback object which is called when received data.
        2 arguments:
        @ ws: the WebSocketApp object
        @ message: utf-8 data received from the server
        """
        # 对收到的message进行解析
        # result = eval(message)
        print("55555555555555")
        print(message)

    def on_error(self,ws, error):

        global reconnect_count
        reconnect_count = 0
        print("-------出错了-----------------")
        print(type(error))
        print(error)
        if type(error) == ConnectionRefusedError or type(
            error) == websocket._exceptions.WebSocketConnectionClosedException:
            print("正在尝试第%d次重连" % reconnect_count)
            reconnect_count += 1
            if reconnect_count < 100:
                self.connection_tmp(ws)
        else:
            print("其他error!")
            print("正在尝试第%d次重连" % reconnect_count)
            reconnect_count += 1
            if reconnect_count < 100:
                self.connection_tmp(ws)

    def on_close(self, ws, close_status_code, close_msg):
        """
        Callback object which is called when the connection is closed.
        2 arguments:
        @ ws: the WebSocketApp object
        @ close_status_code
        @ close_msg
        """
        print('The connection is closed!')

    def start(self):
        # websocket.enableTrace(True)
        self.ws = websocket.WebSocketApp(
            self.url,
            on_open=self.on_open,
            on_message=self.on_message,
            on_data=self.on_data,
            on_error=self.on_error,
            on_close=self.on_close,
        )
        print(f"---------------{self.ws}-----")
        self.ws.run_forever(ping_timeout=3)

    def connection_tmp(self,ws):
        print("进行了一次重连================")
        # websocket.enableTrace(True)
        ws = websocket.WebSocketApp("ws://127.0.0.1:8087",
                                    on_message=self.on_message,
                                    #   on_data=on_data_test,
                                    on_error=self.on_error,
                                    on_close=self.on_close)

        ws.on_open =self.on_open
        try:
            ws.run_forever()
        except KeyboardInterrupt:
            ws.close()
        except:
            ws.close()


if __name__ == "__main__":
    feed = Feed()
    feed.start()

4、参数介绍

（1）url: websocket的地址。

（2）header: 客户发送websocket握手请求的请求头，{'head1:value1','head2:value2'}。

（3）on_open：在建立Websocket握手时调用的可调用对象，这个方法只有一个参数，就是该类本身。

（4）on_message：这个对象在接收到服务器返回的消息时调用。有两个参数，一个是该类本身，一个是我们从服务器获取的字符串（utf-8格式）。

（5）on_error：这个对象在遇到错误时调用，有两个参数，第一个是该类本身，第二个是异常对象。

（6）on_close：在遇到连接关闭的情况时调用，参数只有一个，就是该类本身。

（7）on_cont_message：这个对象在接收到连续帧数据时被调用，有三个参数，分别是：类本身，从服务器接受的字符串（utf-8），连续标志。

（8）on_data：当从服务器接收到消息时被调用，有四个参数，分别是：该类本身，接收到的字符串（utf-8），数据类型，连续标志。

（9）keep_running：一个二进制的标志位，如果为True，这个app的主循环将持续运行，默认值为True。

（10）get_mask_key：用于产生一个掩码。

（11）subprotocols：一组可用的子协议，默认为空。

长连接关键方法：ws.run_forever(ping_interval=60,ping_timeout=5)

如果不断开关闭websocket连接，会一直阻塞下去。另外这个函数带两个参数，如果传的话，启动心跳包发送。

ping_interval:自动发送“ping”命令，每个指定的时间(秒),如果设置为0，则不会自动发送。

ping_timeout:如果没有收到pong消息，则为超时(秒)。

原文链接：https://blog.csdn.net/qq_45664055/article/details/120278303