zwf
131 lines
3.9 KiB
Python
131 lines
3.9 KiB
Python
from abc import ABC
|
|
from typing import Optional, Awaitable, Any, Type
|
|
|
|
from tornado import websocket
|
|
|
|
import tornado.websocket
|
|
|
|
from paste.db.basemodel import BaseModel
|
|
from paste.web.handler import init_user_class
|
|
|
|
|
|
class WebSocketHandler(tornado.websocket.WebSocketHandler, ABC):
|
|
"""
|
|
WebSocketHandler 的派生父类，主要增加了 send 方法，用于向客户端发送数据。
|
|
"""
|
|
|
|
_web_sockets: set['WebSocketHandler'] = set()
|
|
"""
|
|
用于全局保存所有的客户端连接。
|
|
"""
|
|
|
|
user_class: Type[BaseModel] = init_user_class()
|
|
"""
|
|
用户数据处理类。装饰器 web.decorators.auth_token 执行令牌验证时调用该类，用于创建用户对象，并保存在 current_user 属性中。
|
|
注意：这里仅初始化类，而不创建对象。该类允许用户继承扩展，然后自行配置。主要用于执行有关用户的数据操作。
|
|
"""
|
|
|
|
@classmethod
|
|
def add_socket(cls, web_socket):
|
|
"""
|
|
加入 WebSocket 集合。
|
|
|
|
:param web_socket: 要加入的 WebSocketHandler 对象
|
|
"""
|
|
assert hasattr(web_socket, 'send')
|
|
cls._web_sockets.add(web_socket)
|
|
|
|
@classmethod
|
|
def get_sockets(cls):
|
|
"""
|
|
取得 WebSocket 集合。
|
|
"""
|
|
return cls._web_sockets
|
|
|
|
@classmethod
|
|
def has_sockets(cls):
|
|
"""
|
|
连接队列中是否还有连接。
|
|
"""
|
|
return True if cls._web_sockets else False
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
self.token_payload: dict[str: Any] = {}
|
|
"""
|
|
令牌配载数据字典。装饰器 web.decorators.auth_token 执行令牌验证时解码并赋值。 在 HandlerRequest 子类中
|
|
只要配置 auth_token 装饰即可使用该配载数据。
|
|
|
|
其结构为::
|
|
|
|
{
|
|
'iss': private_iss,
|
|
'iat': datetime.datetime.utcnow(),
|
|
'exp': datetime.datetime.utcnow() + datetime.timedelta(days=7),
|
|
'params': {
|
|
'id': user_id,
|
|
'username': username
|
|
}
|
|
}
|
|
"""
|
|
|
|
def token_params(self) -> dict:
|
|
"""
|
|
取出 Token 中的参数字典。
|
|
|
|
:return: 参数字典
|
|
"""
|
|
return self.token_payload.get('params', {})
|
|
|
|
def token_param(self, key):
|
|
"""
|
|
取出 Token 参数字典中的参数。
|
|
|
|
:param key: 参数名称
|
|
"""
|
|
return self.token_params().get(key, None)
|
|
|
|
def is_connected(self):
|
|
"""
|
|
检查当前WebSocket连接是否打开。
|
|
"""
|
|
return self.ws_connection is not None and self.ws_connection.stream is not None
|
|
|
|
def select_subprotocol(self, subprotocols: [str]) -> Optional[str]:
|
|
"""
|
|
选择子协议字符串。注意::
|
|
|
|
1、该方法返回的数据必须位于 subprotocols 数组中；
|
|
2、若有 subprotocols 参数传入，默认始终返回第 0 项；
|
|
3、用于验证的 Token 始终放在子协议的最后一项，读取该数据设置到 request.headers 中；
|
|
|
|
:param subprotocols: 子协议数组，当前端传入字符串时，该数组仅有一项
|
|
:return: 选择的子协议
|
|
"""
|
|
if subprotocols:
|
|
_token = subprotocols[-1]
|
|
self.request.headers.add('Access-Token', _token)
|
|
|
|
return subprotocols[0]
|
|
return None
|
|
|
|
def on_close(self):
|
|
"""
|
|
关闭连接时，从集合中删除客户端连接。
|
|
"""
|
|
if self in self._web_sockets:
|
|
self._web_sockets.remove(self)
|
|
|
|
def send(self) -> Optional[Awaitable[None]]:
|
|
"""
|
|
向客户端发送数据。必须在子类中加以实现。
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
async def data_received(self, chunk: bytes):
|
|
pass
|
|
|
|
def check_origin(self, origin):
|
|
return True
|