feat(client): refactor PushDeliver to support DI & adapt image parsing to the new protocol (#16)

* feat(utils): change the request method of sign to POST

* refactor(client): refactor `PushDeliver` to support dependency injection of Client

* fix(client): fixed parsing images in the new protocol
- ref: LagrangeDev/Lagrange.Core@ef6b41b

Author: pk5ls20

lagrange/__init__.py

@@ -2,6 +2,8 @@
 import asyncio
 
 from .client.client import Client as Client
+from .client.server_push.msg import msg_push_handler
+from .client.server_push.service import server_kick_handler
 from .utils.log import log as log
 from .utils.log import install_loguru as install_loguru
 from .utils.sign import sign_provider
@@ -49,6 +51,8 @@ async def run(self):
             )
             for event, handler in self.events.items():
                 self.client.events.subscribe(event, handler)
+            self.client.push_deliver.subscribe("trpc.msg.olpush.OlPushService.MsgPush", msg_push_handler)
+            self.client.push_deliver.subscribe("trpc.qq_new_tech.status_svc.StatusService.KickNT", server_kick_handler)
             self.client.connect()
             status = await self.login(self.client)
         if not status:

lagrange/client/client.py

@@ -1,7 +1,8 @@
 import os
 import struct
+import asyncio
 from io import BytesIO
-from typing import BinaryIO, Callable, Coroutine, List, Optional, Union, overload
+from typing import BinaryIO, Callable, Coroutine, List, Optional, Union, overload, Literal
 
 from lagrange.info import AppInfo, DeviceInfo, SigInfo
 from lagrange.pb.message.msg_push import MsgPushBody
@@ -32,6 +33,7 @@
     PBGetInfoFromUidReq,
 )
 from lagrange.pb.service.oidb import OidbRequest, OidbResponse
+from lagrange.pb.highway.comm import IndexNode
 from lagrange.utils.binary.protobuf import proto_decode, proto_encode
 from lagrange.utils.log import log
 from lagrange.utils.operator import timestamp
@@ -48,7 +50,7 @@
 from .message.encoder import build_message
 from .message.types import Element
 from .models import UserInfo
-from .server_push import push_handler
+from .server_push.binder import PushDeliver
 from .wtlogin.sso import SSOPacket
 
 
@@ -65,12 +67,17 @@ def __init__(
         super().__init__(uin, app_info, device_info, sig_info, sign_provider, use_ipv6)
 
         self._events = Events()
+        self._push_deliver = PushDeliver(self)
         self._highway = HighWaySession(self)
 
     @property
     def events(self) -> Events:
         return self._events
 
+    @property
+    def push_deliver(self) -> PushDeliver:
+        return self._push_deliver
+
     async def register(self) -> bool:
         if await super().register():
             self._events.emit(ClientOnline(), self)
@@ -155,8 +162,7 @@ async def send_oidb_svc(
         return rsp
 
     async def push_handler(self, sso: SSOPacket):
-        rsp = await push_handler.execute(sso.cmd, sso)
-        if rsp:
+        if rsp := await self._push_deliver.execute(sso.cmd, sso):
             self._events.emit(rsp, self)
 
     async def _send_msg_raw(self, pb: dict, *, grp_id=0, uid="") -> SendMsgRsp:
@@ -229,6 +235,20 @@ async def down_grp_audio(self, audio: Audio, grp_id: int) -> BytesIO:
     async def down_friend_audio(self, audio: Audio) -> BytesIO:
         return await self._highway.download_audio(audio, uid=self.uid)
 
+    async def fetch_image_url(self, bus_type: Literal[10, 20], node: "IndexNode", uid=None, gid=None):
+        if bus_type == 10:
+            return await self._get_pri_img_url(uid, node)
+        elif bus_type == 20:
+            return await self._get_grp_img_url(gid, node)
+        else:
+            raise ValueError("bus_type must be 10 or 20")
+
+    async def _get_grp_img_url(self, grp_id: int, node: "IndexNode") -> str:
+        return await self._highway.get_grp_img_url(grp_id=grp_id, node=node)
+
+    async def _get_pri_img_url(self, uid: str, node: "IndexNode") -> str:
+        return await self._highway.get_pri_img_url(uid=uid, node=node)
+
     async def get_grp_list(self) -> GetGrpListResponse:
         rsp = await self.send_oidb_svc(0xFE5, 2, PBGetGrpListRequest.build().encode())
         if rsp.ret_code:
@@ -292,7 +312,9 @@ async def get_grp_msg(
             and payload.end_seq == end
         ), "return args not matched"
 
-        rsp = [parse_grp_msg(MsgPushBody.decode(i)) for i in payload.elems]
+        rsp = list(
+            await asyncio.gather(*[parse_grp_msg(self, MsgPushBody.decode(i)) for i in payload.elems])
+        )
         if filter_deleted_msg:
             return [*filter(lambda msg: msg.rand != -1, rsp)]
         return rsp

lagrange/client/highway/encoders.py

@@ -225,6 +225,35 @@ def encode_audio_down_req(uuid: str, grp_id: int, uid: str):
     )
 
 
+def encode_grp_img_download_req(grp_id: int, node: IndexNode) -> NTV2RichMediaReq:
+    return NTV2RichMediaReq(
+        req_head=MultiMediaReqHead(
+            common=CommonHead(cmd=200),
+            scene=SceneInfo(
+                req_type=2,
+                bus_type=1,
+                scene_type=2,
+                grp=GroupInfo(grp_id=grp_id),
+            )
+        ),
+        download=DownloadReq(node=node),
+    )
+
+
+def encode_pri_img_download_req(uid: str, node: IndexNode) -> NTV2RichMediaReq:
+    return NTV2RichMediaReq(
+        req_head=MultiMediaReqHead(
+            common=CommonHead(cmd=200),
+            scene=SceneInfo(
+                req_type=2,
+                bus_type=1,
+                scene_type=1,
+                c2c=C2CUserInfo(uid=uid),
+            )
+        ),
+        download=DownloadReq(node=node),
+    )
+
 # def encode_video_upload_req(
 #         seq: int,
 #         from_uin: int,

lagrange/client/highway/highway.py

@@ -6,6 +6,7 @@
 from typing import TYPE_CHECKING, BinaryIO, List, Optional, Tuple
 
 from lagrange.client.message.elems import Audio, Image
+from lagrange.pb.highway.comm import IndexNode
 from lagrange.pb.highway.ext import NTV2RichMediaHighwayExt
 from lagrange.pb.highway.httpconn import HttpConn0x6ffReq, HttpConn0x6ffRsp
 from lagrange.pb.highway.rsp import NTV2RichMediaResp, DownloadInfo, DownloadRsp
@@ -21,6 +22,8 @@
     encode_highway_head,
     encode_upload_img_req,
     encode_audio_down_req,
+    encode_grp_img_download_req,
+    encode_pri_img_download_req,
 )
 from .frame import read_frame, write_frame
 from .utils import calc_file_hash_and_length, timeit
@@ -240,6 +243,38 @@ async def upload_image(self, file: BinaryIO, gid=0, uid="") -> Image:
             qmsg=None if gid else ret.upload.compat_qmsg,
         )
 
+    async def get_grp_img_url(self, grp_id: int, node: "IndexNode") -> str:
+        ret = NTV2RichMediaResp.decode(
+            (
+                await self._client.send_oidb_svc(
+                    0x11C4,
+                    200,
+                    encode_grp_img_download_req(
+                        grp_id, node
+                    ).encode(),
+                    True
+                )
+            ).data
+        )
+        body = ret.download
+        return f"https://{body.info.domain}{body.info.url_path}{body.rkey}"
+
+    async def get_pri_img_url(self, uid: str, node: IndexNode) -> str:
+        ret = NTV2RichMediaResp.decode(
+            (
+                await self._client.send_oidb_svc(
+                    0x11C5,
+                    200,
+                    encode_pri_img_download_req(
+                        uid, node
+                    ).encode(),
+                    True
+                )
+            ).data
+        )
+        body = ret.download
+        return f"https://{body.info.domain}{body.info.url_path}{body.rkey}"
+
     async def upload_voice(self, file: BinaryIO, gid=0, uid="") -> Audio:
         if not self._session_addr_list:
             await self._get_bdh_session()

lagrange/client/message/decoder.py

@@ -1,5 +1,5 @@
 import zlib
-from typing import List, Tuple, Sequence
+from typing import List, Tuple, Sequence, TYPE_CHECKING, cast, Literal
 
 from lagrange.client.events.group import GroupMessage
 from lagrange.client.events.friend import FriendMessage
@@ -9,6 +9,10 @@
 from . import elems
 from .types import Element
 from lagrange.utils.binary.protobuf import proto_encode
+from lagrange.pb.highway.comm import MsgInfo
+
+if TYPE_CHECKING:
+    from lagrange.client.client import Client
 
 
 def parse_msg_info(pb: MsgPushBody) -> Tuple[int, str, int, int, int]:
@@ -31,7 +35,8 @@ def parse_friend_info(pkg: MsgPushBody) -> Tuple[int, str, int, str]:
     return from_uin, from_uid, to_uin, to_uid
 
 
-def parse_msg_new(rich: RichText) -> Sequence[Element]:
+async def parse_msg_new(client: "Client", pkg: MsgPushBody) -> Sequence[Element]:
+    rich: RichText = pkg.message.body
     if rich.ptt:
         ptt = rich.ptt
         return [
@@ -134,6 +139,27 @@ def parse_msg_new(rich: RichText) -> Sequence[Element]:
                         f8=common.pb_elem[8],
                     )
                 )
+            if common.bus_type in [10, 20]:  # 10: friend, 20: group
+                extra = MsgInfo.decode(proto_encode(raw.common_elem.pb_elem))
+                index = extra.body[0].index
+                uid = client.uid
+                gid = pkg.response_head.rsp_grp.gid if common.bus_type == 20 else None
+                url = await client.fetch_image_url(bus_type=cast(Literal[10, 20], common.bus_type),
+                                                   node=index, uid=uid, gid=gid)
+                msg_chain.append(
+                    elems.Image(
+                        name=index.info.name,
+                        size=index.info.size,
+                        id=0,
+                        md5=bytes.fromhex(index.info.hash),
+                        text=extra.biz_info.pic.summary if extra.biz_info.pic.summary else "[图片]",
+                        width=index.info.width,
+                        height=index.info.height,
+                        url=url,
+                        is_emoji=extra.biz_info.pic.biz_type != 0,
+                        qmsg=None,
+                    )
+                )
         elif raw.rich_msg:
             service = raw.rich_msg
             if service.template:
@@ -227,13 +253,13 @@ def parse_msg_new(rich: RichText) -> Sequence[Element]:
     return msg_chain
 
 
-def parse_friend_msg(pkg: MsgPushBody) -> FriendMessage:
+async def parse_friend_msg(client: "Client", pkg: MsgPushBody) -> FriendMessage:
     from_uin, from_uid, to_uin, to_uid = parse_friend_info(pkg)
 
     seq = pkg.content_head.seq
     msg_id = pkg.content_head.msg_id
     timestamp = pkg.content_head.timestamp
-    parsed_msg = parse_msg_new(pkg.message.body)
+    parsed_msg = await parse_msg_new(client, pkg)
     msg_text = "".join([getattr(msg, "text", "") for msg in parsed_msg])
 
     return FriendMessage(
@@ -249,7 +275,7 @@ def parse_friend_msg(pkg: MsgPushBody) -> FriendMessage:
     )
 
 
-def parse_grp_msg(pkg: MsgPushBody) -> GroupMessage:
+async def parse_grp_msg(client: "Client", pkg: MsgPushBody) -> GroupMessage:
     user_id, uid, seq, time, rand = parse_msg_info(pkg)
 
     grp_id = pkg.response_head.rsp_grp.gid
@@ -261,7 +287,7 @@ def parse_grp_msg(pkg: MsgPushBody) -> GroupMessage:
     if isinstance(grp_name, bytes):  # unexpected end of data
         grp_name = grp_name.decode("utf-8", errors="ignore")
 
-    parsed_msg = parse_msg_new(pkg.message.body)
+    parsed_msg = await parse_msg_new(client, pkg)
     msg_text = "".join([getattr(msg, "text", "") for msg in parsed_msg])
 
     return GroupMessage(

lagrange/client/server_push/__init__.py

@@ -1,4 +0,0 @@
-from . import msg, service
-from .binder import push_handler
-
-__all__ = ["push_handler"]

lagrange/client/server_push/binder.py

@@ -1,35 +1,25 @@
-import functools
-from typing import Any, Callable, Coroutine, Dict
+from typing import Any, Callable, Coroutine, Dict, TYPE_CHECKING
 
 from lagrange.client.wtlogin.sso import SSOPacket
 
 from .log import logger
 
+if TYPE_CHECKING:
+    from lagrange.client.client import Client
+
 
 class PushDeliver:
-    def __init__(self):
+    def __init__(self, client: "Client"):
+        self._client = client
         self._handle_map: Dict[
-            str, Callable[[SSOPacket], Coroutine[None, None, Any]]
+            str, Callable[["Client", SSOPacket], Coroutine[None, None, Any]]
         ] = {}
 
-    def subscribe(self, cmd: str):
-        def _decorator(
-            func: Callable[[SSOPacket], Coroutine[None, None, Any]]
-        ) -> Callable[[SSOPacket], Coroutine[None, None, Any]]:
-            @functools.wraps(func)
-            async def _wrapper(packet: SSOPacket):
-                return await func(packet)
-
-            self._handle_map[cmd] = _wrapper  # noqa
-            return _wrapper
-
-        return _decorator
+    def subscribe(self, cmd: str, func: Callable[["Client", SSOPacket], Coroutine[None, None, Any]]):
+        self._handle_map[cmd] = func
 
     async def execute(self, cmd: str, sso: SSOPacket):
         if cmd not in self._handle_map:
-            logger.warning("unsupported command: {}".format(cmd))
+            logger.warning(f"Unsupported command: {cmd}")
         else:
-            return await self._handle_map[cmd](sso)
-
-
-push_handler = PushDeliver()
+            return await self._handle_map[cmd](self._client, sso)

lagrange/client/server_push/msg.py

@@ -1,5 +1,6 @@
 import json
 import re
+from typing import TYPE_CHECKING
 
 from lagrange.client.message.decoder import parse_grp_msg, parse_friend_msg
 from lagrange.pb.message.msg_push import MsgPush
@@ -26,21 +27,22 @@
     GroupRecall,
 )
 from ..wtlogin.sso import SSOPacket
-from .binder import push_handler
 from .log import logger
 
+if TYPE_CHECKING:
+    from lagrange.client.client import Client
 
-@push_handler.subscribe("trpc.msg.olpush.OlPushService.MsgPush")
-async def msg_push_handler(sso: SSOPacket):
+
+async def msg_push_handler(client: "Client", sso: SSOPacket):
     pkg = MsgPush.decode(sso.data).body
     typ = pkg.content_head.type
     sub_typ = pkg.content_head.sub_type
 
     logger.debug("msg_push received, type: {}.{}".format(typ, sub_typ))
     if typ == 82:  # grp msg
-        return parse_grp_msg(pkg)
+        return await parse_grp_msg(client, pkg)
     elif typ == 166:  # frd msg
-        return parse_friend_msg(pkg)
+        return await parse_friend_msg(client, pkg)
     elif typ == 33:  # member joined
         pb = MemberChanged.decode(pkg.message.buf2)
         return GroupMemberJoined(

lagrange/client/server_push/service.py

@@ -2,10 +2,8 @@
 
 from ..events.service import ServerKick
 from ..wtlogin.sso import SSOPacket
-from .binder import push_handler
 
 
-@push_handler.subscribe("trpc.qq_new_tech.status_svc.StatusService.KickNT")
-async def server_kick(sso: SSOPacket):
+async def server_kick_handler(_, sso: SSOPacket):
     ev = KickNT.decode(sso.data)
     return ServerKick(tips=ev.tips, title=ev.title)