feat: websocket

aiotieba/__init__.py

@@ -33,6 +33,7 @@
     TypeFragTiebaPlus,
     UserInfo_home,
 )
+from .client.websocket import Websocket
 
 if os.name == 'posix':
     import signal

aiotieba/__version__.py

@@ -1 +1 @@
-__version__ = "3.1.5a2"
+__version__ = "3.1.5a3"

aiotieba/client/_client.py

@@ -11,8 +11,8 @@
 from ._core import TbCore
 from ._helper import ForumInfoCache, is_portrait
 from ._typing import TypeUserInfo
-from ._websocket import Websocket
 from .get_homepage._classdef import UserInfo_home
+from .websocket import Websocket
 
 if TYPE_CHECKING:
     import numpy as np
@@ -75,6 +75,9 @@ def __init__(
         loop: Optional[asyncio.AbstractEventLoop] = None,
     ) -> None:
 
+        if loop is None:
+            loop = asyncio.get_running_loop()
+
         self._core = TbCore(BDUSS_key, proxy, loop)
         self._user = UserInfo_home()._init_null()
 
@@ -2059,13 +2062,12 @@ async def send_msg(self, _id: Union[str, int], content: str) -> bool:
 
         try:
             await self.websocket.init_websocket()
-            self.websocket._record_id += 1
 
             from . import send_msg
 
-            proto = send_msg.pack_proto(user_id, content, self.websocket._record_id)
-            body = await self.websocket.send(proto, send_msg.CMD, timeout=5.0)
-            send_msg.parse_body(body)
+            proto = send_msg.pack_proto(user_id, content, self.websocket.record_id)
+            resp = await self.websocket.send(proto, send_msg.CMD)
+            send_msg.parse_body(await resp.read())
 
         except Exception as err:
             LOG().warning(f"{err}. user_id={user_id}")

aiotieba/client/_helper/__init__.py

@@ -11,7 +11,9 @@
     pack_proto_request,
     pack_web_form_request,
     pack_web_get_request,
+    pack_ws_bytes,
     parse_json,
+    parse_ws_bytes,
     removeprefix,
     removesuffix,
     send_request,

aiotieba/client/_helper/_func.py

@@ -271,6 +271,75 @@ def pack_web_form_request(core: TbCore, url: yarl.URL, data: List[Tuple[str, str
     return request
 
 
+def pack_ws_bytes(
+    core: TbCore, data: bytes, cmd: int, req_id: int, *, compress: bool = False, encrypt: bool = True
+) -> bytes:
+    """
+    打包数据并添加9字节头部
+
+    Args:
+        core (TiebaCore): 贴吧核心参数容器
+        data (bytes): 待发送的websocket数据
+        cmd (int): 请求的cmd类型
+        req_id (int): 请求的id
+        compress (bool, optional): 是否需要gzip压缩. Defaults to False.
+        encrypt (bool, optional): 是否需要aes加密. Defaults to True.
+
+    Returns:
+        bytes: 打包后的websocket数据
+    """
+
+    flag = 0x08
+
+    if compress:
+        flag |= 0b01000000
+        data = gzip.compress(data, compresslevel=-1, mtime=0)
+    if encrypt:
+        flag |= 0b10000000
+        data = pad(data, AES.block_size)
+        data = core.aes_ecb_chiper.encrypt(data)
+
+    data = b''.join(
+        [
+            flag.to_bytes(1, 'big'),
+            cmd.to_bytes(4, 'big'),
+            req_id.to_bytes(4, 'big'),
+            data,
+        ]
+    )
+
+    return data
+
+
+def parse_ws_bytes(core: TbCore, data: bytes) -> Tuple[bytes, int, int]:
+    """
+    对websocket返回数据进行解包
+
+    Args:
+        core (TiebaCore): 贴吧核心参数容器
+        data (bytes): 接收到的websocket数据
+
+    Returns:
+        bytes: 解包后的websocket数据
+        int: 对应请求的cmd类型
+        int: 对应请求的id
+    """
+
+    data_view = memoryview(data)
+    flag = data_view[0]
+    cmd = int.from_bytes(data_view[1:5], 'big')
+    req_id = int.from_bytes(data_view[5:9], 'big')
+
+    data = data_view[9:].tobytes()
+    if flag & 0b10000000:
+        data = core.aes_ecb_chiper.decrypt(data)
+        data = unpad(data, AES.block_size)
+    if flag & 0b01000000:
+        data = gzip.decompress(data)
+
+    return data, cmd, req_id
+
+
 def check_status_code(response: aiohttp.ClientResponse) -> None:
     if response.status != 200:
         raise HTTPStatusError(response.status, response.reason)

aiotieba/client/get_group_msg/__init__.py

@@ -1,2 +1,2 @@
 from ._api import CMD, pack_proto, parse_body
-from ._classdef import MsgGroup
+from ._classdef import GroupType

aiotieba/client/get_group_msg/_api.py

@@ -2,27 +2,36 @@
 
 from .._core import TbCore
 from .._exception import TiebaServerError
-from ._classdef import MsgGroup
+from ..push_notify import WsNotify
+from ._classdef import WsMsgGroup
 from .protobuf import GetGroupMsgReqIdl_pb2, GetGroupMsgResIdl_pb2
 
 CMD = 202003
 
 
-def pack_proto(core: TbCore, groups: List[MsgGroup], get_type: int) -> bytes:
+def pack_proto(core: TbCore, notifies: List[WsNotify], get_type: int) -> bytes:
     req_proto = GetGroupMsgReqIdl_pb2.GetGroupMsgReqIdl()
-    for group in groups:
+    req_proto.data.width = 720
+    req_proto.data.height = 1280
+    req_proto.data.smallWidth = 240
+    req_proto.data.smallHeight = 240
+    for notify in notifies:
         group_proto = req_proto.data.groupMids.add()
-        group_proto.groupId = group._group_type
-        group_proto.lastMsgId = group._last_msg_id
+        group_proto.groupId = notify._group_id
+        group_proto.lastMsgId = notify._last_msg_id
     req_proto.data.gettype = str(get_type)
     req_proto.cuid = f"{core.cuid}|com.baidu.tieba_mini{core.post_version}"
 
     return req_proto.SerializeToString()
 
 
-def parse_body(body: bytes) -> None:
+def parse_body(body: bytes) -> List[WsMsgGroup]:
     res_proto = GetGroupMsgResIdl_pb2.GetGroupMsgResIdl()
     res_proto.ParseFromString(body)
 
     if code := res_proto.error.errorno:
         raise TiebaServerError(code, res_proto.error.errmsg)
+
+    groups = [WsMsgGroup()._init(p) for p in res_proto.data.groupInfo]
+
+    return groups