python 使用的WebSocket和bearychat进行长连接。当服务器接收到消息后执行动作返回消息

Posted 2021-05-09

import re
import json
from urllib.parse import urlencode
from config import BEARYCHAT_TOKEN
from tornado.ioloop import IOLoop, PeriodicCallback
from tornado.httpclient import AsyncHTTPClient
from tornado import gen
from tornado.websocket import websocket_connect


class Handler():
    handle_register = {}

    def __init__(self, func):
        Handler.handle_register[func.__name__] = func

    @classmethod
    def get(cls, k):
        for i in cls.handle_register.keys():
            if re.search(i, k):
                return cls.handle_register.get(i)
        return False

    def __call__(self, func):
        return func


@Handler
def test():
    return "Hello,World"


class Client(object):
    def __init__(self):
        self.ioloop = IOLoop.instance()
        self.ws = None
        self.connect()
        PeriodicCallback(self.keep_alive, 20000, io_loop=self.ioloop).start()
        self.ioloop.start()

    @gen.coroutine
    def connect(self):
        print("trying to connect")
        try:
            post_data = {'token': BEARYCHAT_TOKEN}
            get_url = yield AsyncHTTPClient().fetch("https://rtm.bearychat.com/start", method="POST",
                                                    body=urlencode(post_data), connect_timeout=5, request_timeout=5)
            url = json.loads(get_url.body.decode()).get('result', {}).get('ws_host')
            self.ws = yield websocket_connect(url, connect_timeout=5)
        except Exception as e:
            print("connection error,{}".format(e))
        else:
            print("connected")
            self.run()

    @gen.coroutine
    def run(self):
        while True:
            msg = yield self.ws.read_message()
            if msg is None:
                self.ws = None
                break
            try:
                msg = json.loads(msg)
            except:
                self.ws = None
                break
            if msg.get('type') != "channel_message":
                continue
            print(msg)
            replay_func = Handler.get(msg.get("text"))
            if replay_func:
                raw_text = replay_func()
                finish_text = json.dumps(
                    {"text": raw_text,
                     "vchannel_id": msg.get("vchannel_id"),
                     "call_id": 23,
                     "refer_key": msg.get("key"),
                     "type": "channel_message",
                     "channel_id": msg.get("channel_id")}
                )
                self.ws.write_message(finish_text)

    def keep_alive(self):
        if self.ws is None:
            self.connect()
        else:
            self.ws.write_message('{"call_id": 29, "type": "ping"}')


if __name__ == "__main__":
    Client()