satori-python-adapter-milky 0.1.0__py3-none-any.whl

satori/adapters/milky/main.py

@@ -0,0 +1,247 @@
+from __future__ import annotations
+
+import asyncio
+from datetime import datetime
+
+import aiohttp
+from launart import Launart, any_completed
+from launart.status import Phase
+from loguru import logger
+from starlette.responses import JSONResponse, Response
+from yarl import URL
+
+from satori import EventType
+from satori.exception import ActionFailed
+from satori.model import Event, Login, LoginStatus
+from satori.server.adapter import Adapter as BaseAdapter
+from satori.server.model import Request
+from satori.utils import decode
+
+from .api import apply
+from .events import event_handlers
+from .utils import MilkyNetwork, decode_login_user
+
+DEFAULT_FEATURES = ["guild.plain", "reaction"]
+
+
+class _MilkyNetwork:
+    def __init__(self, adapter: MilkyAdapter):  # type: ignore[name-defined]
+        self.adapter = adapter
+
+    async def call_api(self, action: str, params: dict | None = None):
+        return await self.adapter.call_api(action, params or {})
+
+
+class MilkyAdapter(BaseAdapter):
+
+    session: aiohttp.ClientSession | None
+    connection: aiohttp.ClientWebSocketResponse | None
+
+    def __init__(
+        self,
+        endpoint: str | URL,
+        *,
+        token: str | None = None,
+        token_in_query: bool = False,
+        headers: dict[str, str] | None = None,
+    ):
+        super().__init__()
+        self.base_url = URL(str(endpoint))
+        base_path = self.base_url.path.rstrip("/")
+        self.api_base = self.base_url.with_path(f"{base_path}/api")
+        ws_scheme = "wss" if self.base_url.scheme == "https" else "ws"
+        self.event_url = self.base_url.with_scheme(ws_scheme).with_path(f"{base_path}/event")
+        if token_in_query and token:
+            self.event_url = self.event_url.update_query(access_token=token)
+        self.token = token
+        self.headers = headers.copy() if headers else {}
+        self.session = None
+        self.connection = None
+        self.close_signal = asyncio.Event()
+        self.logins: dict[str, Login] = {}
+        self.networks: dict[str, MilkyNetwork] = {}
+        self.features = list(DEFAULT_FEATURES)
+        apply(self, self._get_network, self._get_login)
+
+    def get_platform(self) -> str:
+        return "milky"
+
+    def ensure(self, platform: str, self_id: str) -> bool:
+        return platform == "milky" and self_id in self.logins
+
+    async def get_logins(self) -> list[Login]:
+        logins = list(self.logins.values())
+        for index, login in enumerate(logins):
+            login.sn = index
+        return logins
+
+    @property
+    def required(self) -> set[str]:
+        return {"satori-python.server"}
+
+    @property
+    def stages(self) -> set[Phase]:
+        return {"preparing", "blocking", "cleanup"}
+
+    async def launch(self, manager: Launart):
+        async with self.stage("preparing"):
+            self.session = aiohttp.ClientSession()
+
+        async with self.stage("blocking"):
+            await self.connection_daemon(manager, self.session)
+
+        async with self.stage("cleanup"):
+            if self.connection and not self.connection.closed:
+                await self.connection.close()
+            if self.session:
+                await self.session.close()
+            self.connection = None
+            self.session = None
+            await self._handle_disconnect()
+
+    def proxy_urls(self) -> list[str]:
+        return []
+
+    async def handle_internal(self, request: Request, path: str) -> Response:
+        if path.startswith("_api"):
+            data = await request.origin.json()
+            return JSONResponse(await self.call_api(path[5:], data))
+        if not self.session:
+            raise RuntimeError("HTTP session not initialized")
+        url = self.base_url.with_path(path)
+        headers = self.headers.copy()
+        if self.token:
+            headers.setdefault("Authorization", f"Bearer {self.token}")
+        async with self.session.get(url, headers=headers) as resp:
+            content = await resp.read()
+            return Response(content=content, media_type=resp.headers.get("Content-Type"))
+
+    async def call_api(self, action: str, params: dict | None = None) -> dict | None:
+        if not self.session:
+            raise RuntimeError("HTTP session not initialized")
+        url = self.api_base.with_path(f"{self.api_base.path.rstrip('/')}/{action}")
+        headers = self.headers.copy()
+        if self.token:
+            headers.setdefault("Authorization", f"Bearer {self.token}")
+        async with self.session.post(url, json=params or {}, headers=headers) as resp:
+            resp.raise_for_status()
+            data = await resp.json()
+        if data.get("status") == "failed":
+            raise ActionFailed(f"{data.get('retcode')}: {data.get('message')}", data)
+        return data.get("data")
+
+    async def connection_daemon(self, manager: Launart, session: aiohttp.ClientSession):
+        while not manager.status.exiting:
+            headers = self.headers.copy()
+            if self.token:
+                headers.setdefault("Authorization", f"Bearer {self.token}")
+            try:
+                self.connection = await session.ws_connect(self.event_url, headers=headers)
+            except Exception as e:
+                logger.error(f"Milky adapter websocket connect failed: {e}")
+                await asyncio.sleep(5)
+                continue
+            logger.info("Milky adapter websocket connected")
+            self.close_signal.clear()
+            await self.refresh_login()
+            receiver_task = asyncio.create_task(self.message_handle())
+            close_task = asyncio.create_task(self.close_signal.wait())
+            sigexit_task = asyncio.create_task(manager.status.wait_for_sigexit())
+
+            done, pending = await any_completed(receiver_task, close_task, sigexit_task)
+            for task in pending:
+                task.cancel()
+            await asyncio.gather(*pending, return_exceptions=True)
+            if sigexit_task in done:
+                break
+            logger.warning("Milky adapter websocket closed, retrying in 5 seconds")
+            await self._handle_disconnect()
+            await asyncio.sleep(5)
+        await self._handle_disconnect()
+
+    async def message_handle(self):
+        assert self.connection is not None
+        async for msg in self.connection:
+            if msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED):
+                self.close_signal.set()
+                break
+            if msg.type != aiohttp.WSMsgType.TEXT:
+                continue
+            try:
+                data = decode(msg.data)
+            except Exception as e:  # pragma: no cover - defensive
+                logger.error(f"Failed to decode milky event: {e}")
+                continue
+            if not isinstance(data, dict):
+                continue
+            await self.handle_event(data)
+
+    async def handle_event(self, payload: dict):
+        event_type = payload.get("event_type")
+        self_id = str(payload.get("self_id"))
+        if not event_type or not self_id:
+            return
+        if self_id not in self.logins:
+            await self.refresh_login()
+        if self_id not in self.logins:
+            logger.warning(f"Ignoring event for unknown self_id {self_id}")
+            return
+        login = self.logins[self_id]
+        handler = event_handlers.get(event_type)
+        network = self.networks.get(self_id)
+        if not network:
+            network = _MilkyNetwork(self)
+            self.networks[self_id] = network
+        if handler:
+            event = await handler(login, network, payload)
+        else:
+            body = payload.get("data", {})
+            event = Event(
+                EventType.INTERNAL,
+                datetime.fromtimestamp(payload.get("time", datetime.now().timestamp())),
+                login,
+                _type=event_type,
+                _data=body,
+            )
+        if event:
+            await self.server.post(event)
+
+    async def refresh_login(self):
+        try:
+            data = await self.call_api("get_login_info", {})
+        except Exception as e:
+            logger.error(f"Failed to fetch milky login info: {e}")
+            return
+        if not data:
+            return
+        user = decode_login_user(data)
+        login = Login(0, LoginStatus.ONLINE, "milky", platform="milky", user=user, features=self.features.copy())
+        self_id = login.id
+        previous = self.logins.get(self_id)
+        self.logins[self_id] = login
+        self.networks[self_id] = _MilkyNetwork(self)
+        event_type = EventType.LOGIN_ADDED if previous is None else EventType.LOGIN_UPDATED
+        await self.server.post(Event(event_type, datetime.now(), login))
+
+    async def _handle_disconnect(self):
+        for self_id, login in list(self.logins.items()):
+            login.status = LoginStatus.OFFLINE
+            await self.server.post(Event(EventType.LOGIN_REMOVED, datetime.now(), login))
+            self.logins.pop(self_id, None)
+            self.networks.pop(self_id, None)
+        if self.connection and not self.connection.closed:
+            await self.connection.close()
+        self.close_signal.set()
+
+    def _get_network(self, self_id: str) -> MilkyNetwork:
+        network = self.networks.get(self_id)
+        if not network:
+            network = _MilkyNetwork(self)
+            self.networks[self_id] = network
+        return network
+
+    def _get_login(self, self_id: str) -> Login:
+        return self.logins[self_id]
+
+
+__all__ = ["MilkyAdapter"]

satori/adapters/milky/message.py

@@ -0,0 +1,233 @@
+from __future__ import annotations
+
+import re
+from collections.abc import Sequence
+from datetime import datetime
+from typing import Any
+
+from satori.element import Custom, E, Element
+from satori.model import Channel, ChannelType, Login, MessageObject, User
+from satori.parser import Element as RawElement
+from satori.parser import parse
+
+from .utils import MilkyNetwork, decode_guild, decode_guild_channel_id, decode_member, get_scene_and_peer, user_avatar
+
+_BASE64_RE = re.compile(r"^data:([\w/.+-]+);base64,")
+
+
+class MilkyMessageEncoder:
+    def __init__(self, login: Login, net: MilkyNetwork, channel_id: str):
+        self.login = login
+        self.net = net
+        self.channel_id = channel_id
+        self.segments: list[dict[str, Any]] = []
+        self.elements: list[Element] = []
+        self.results: list[MessageObject] = []
+
+    async def send(self, content: str) -> list[MessageObject]:
+        raw_elements = parse(content)
+        await self.render(raw_elements)
+        await self.flush()
+        return self.results
+
+    async def render(self, elements: Sequence[RawElement]):
+        for element in elements:
+            await self.visit(element)
+
+    async def visit(self, element: RawElement):
+        type_ = element.type
+        attrs = element.attrs
+        children = element.children
+        if type_ == "text":
+            text = attrs.get("text", "")
+            if not self.segments or self.segments[-1]["type"] != "text":
+                self.segments.append({"type": "text", "data": {"text": text}})
+            else:
+                self.segments[-1]["data"]["text"] += text
+            self.elements.append(E.text(text))
+        elif type_ == "br":
+            if not self.segments or self.segments[-1]["type"] != "text":
+                self.segments.append({"type": "text", "data": {"text": "\n"}})
+            else:
+                self.segments[-1]["data"]["text"] += "\n"
+            self.elements.append(E.text("\n"))
+        elif type_ == "p":
+            await self.render(children)
+            if self.segments:
+                if self.segments[-1]["type"] == "text":
+                    self.segments[-1]["data"]["text"] += "\n"
+                else:
+                    self.segments.append({"type": "text", "data": {"text": "\n"}})
+            self.elements.append(E.text("\n"))
+        elif type_ == "at":
+            if attrs.get("type") == "all":
+                self.segments.append({"type": "mention_all", "data": {}})
+                self.elements.append(E.at_all())
+            elif "id" in attrs:
+                target = attrs["id"]
+                self.segments.append({"type": "mention", "data": {"user_id": int(target)}})
+                self.elements.append(E.at(str(target)))
+        elif type_ == "quote":
+            quote_id = attrs.get("id")
+            if quote_id is not None:
+                try:
+                    seq = int(quote_id)
+                except ValueError:
+                    seq = None
+                if seq is not None:
+                    self.segments.append({"type": "reply", "data": {"message_seq": seq}})
+                    self.elements.append(E.quote(str(seq), content=[E.text("")]))
+            await self.render(children)
+        elif type_ in {"img", "image"}:
+            uri = attrs.get("src") or attrs.get("url")
+            if not uri:
+                return
+            if match := _BASE64_RE.match(uri):
+                uri = f"base64://{uri[len(match.group(0)) :]}"
+            self.segments.append({"type": "image", "data": {"uri": uri, "sub_type": attrs.get("sub_type", "normal")}})
+            self.elements.append(E.image(uri))
+        elif type_ == "audio":
+            uri = attrs.get("src") or attrs.get("url")
+            if not uri:
+                return
+            if match := _BASE64_RE.match(uri):
+                uri = f"base64://{uri[len(match.group(0)) :]}"
+            self.segments.append({"type": "record", "data": {"uri": uri}})
+            self.elements.append(E.audio(uri))
+        elif type_ == "video":
+            uri = attrs.get("src") or attrs.get("url")
+            if not uri:
+                return
+            if match := _BASE64_RE.match(uri):
+                uri = f"base64://{uri[len(match.group(0)) :]}"
+            payload = {"uri": uri}
+            if poster := attrs.get("poster"):
+                payload["thumb_uri"] = poster
+            self.segments.append({"type": "video", "data": payload})
+            self.elements.append(E.video(uri))
+        else:
+            await self.render(children)
+
+    async def flush(self):
+        if not self.segments:
+            return
+        scene, peer_id = get_scene_and_peer(self.channel_id)
+        payload: dict = {"message": self.segments}
+        if scene == "group":
+            payload["group_id"] = peer_id
+            resp = await self.net.call_api("send_group_message", payload)
+        else:
+            payload["user_id"] = peer_id
+            resp = await self.net.call_api("send_private_message", payload)
+        if resp:
+            channel_type = ChannelType.TEXT if scene == "group" else ChannelType.DIRECT
+            channel_id = str(peer_id) if scene == "group" else self.channel_id
+            channel = Channel(channel_id, channel_type)
+            created_at = datetime.fromtimestamp(resp.get("time", datetime.now().timestamp()))
+            content = "".join(str(elem) for elem in self.elements)
+            message = MessageObject(str(resp.get("message_seq", "")), content, channel=channel, created_at=created_at)
+            message.message = list(self.elements)
+            self.results.append(message)
+        self.segments = []
+        self.elements = []
+
+
+async def decode_message(net: MilkyNetwork, payload: dict) -> MessageObject:
+    elements = await _decode_segments(net, payload, payload.get("segments") or [])
+    guild_id, channel_id = decode_guild_channel_id(payload)
+    channel_type = ChannelType.TEXT if guild_id else ChannelType.DIRECT
+    channel_name = None
+    guild = None
+    member = None
+    if payload["message_scene"] == "group":
+        group_info = payload.get("group")
+        if group_info:
+            guild = decode_guild(group_info)
+            channel_name = group_info.get("group_name")
+        member_info = payload.get("group_member")
+        if member_info:
+            member = decode_member(member_info)
+    elif payload["message_scene"] == "friend":
+        friend_info = payload.get("friend")
+        if friend_info:
+            channel_name = friend_info.get("nickname")
+    channel = Channel(channel_id, channel_type, channel_name)
+
+    user_name = None
+    if payload["message_scene"] == "group":
+        member_info = payload.get("group_member")
+        if member_info:
+            user_name = member_info.get("nickname")
+    elif payload["message_scene"] == "friend":
+        friend_info = payload.get("friend")
+        if friend_info:
+            user_name = friend_info.get("nickname")
+    user = User(str(payload["sender_id"]), user_name, avatar=user_avatar(payload["sender_id"]))
+
+    message = MessageObject(
+        str(payload["message_seq"]),
+        "".join(str(elem) for elem in elements),
+        channel=channel,
+        guild=guild,
+        member=member,
+        user=user,
+        created_at=datetime.fromtimestamp(payload["time"]),
+    )
+    message.message = elements
+    return message
+
+
+async def _decode_segments(net: MilkyNetwork, payload: dict, segments: Sequence[dict]) -> list[Element]:
+    result: list[Element] = []
+    for segment in segments:
+        seg_type = segment.get("type")
+        data = segment.get("data", {})
+        if seg_type == "text":
+            result.append(E.text(data.get("text", "")))
+        elif seg_type == "mention":
+            result.append(E.at(str(data.get("user_id"))))
+        elif seg_type == "mention_all":
+            result.append(E.at_all())
+        elif seg_type == "image":
+            result.append(E.image(_resource_url(data)))
+        elif seg_type == "record":
+            result.append(E.audio(_resource_url(data)))
+        elif seg_type == "video":
+            result.append(E.video(_resource_url(data)))
+        elif seg_type == "file":
+            result.append(E.file(_resource_url(data)))
+        elif seg_type == "reply":
+            seq = data.get("message_seq")
+            if seq is not None:
+                quote = await _decode_reply(net, payload, int(seq))
+                if quote:
+                    result.append(quote)
+        else:
+            result.append(Custom(f"milky:{seg_type}", data))
+    return result
+
+
+async def _decode_reply(net: MilkyNetwork, payload: dict, message_seq: int) -> Element | None:
+    try:
+        response = await net.call_api(
+            "get_message",
+            {
+                "message_scene": payload["message_scene"],
+                "peer_id": payload["peer_id"],
+                "message_seq": message_seq,
+            },
+        )
+    except Exception:
+        return None
+    if not response or "message" not in response:
+        return None
+    quoted = await decode_message(net, response["message"])
+    content = []
+    if quoted.user:
+        content.append(E.author(quoted.user.id, quoted.user.name, quoted.user.avatar))
+    content.extend(quoted.message)
+    return E.quote(str(message_seq), content=content)
+
+
+def _resource_url(data: dict) -> str:
+    return data.get("temp_url") or data.get("url") or data.get("uri") or ""

satori/adapters/milky/utils.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Literal, Protocol
+
+from satori.model import Channel, ChannelType, Guild, Member, User
+
+AVATAR_URL = "https://q.qlogo.cn/headimg_dl?dst_uin={uin}&spec=640"
+GROUP_AVATAR_URL = "https://p.qlogo.cn/gh/{group}/{group}/640"
+
+
+class MilkyNetwork(Protocol):
+    async def call_api(self, action: str, params: dict | None = None) -> dict | None: ...
+
+
+def user_avatar(uin: int | str) -> str:
+    return AVATAR_URL.format(uin=uin)
+
+
+def group_avatar(group_id: int | str) -> str:
+    return GROUP_AVATAR_URL.format(group=group_id)
+
+
+def decode_group_channel(group: dict) -> Channel:
+    return Channel(str(group["group_id"]), ChannelType.TEXT, group.get("group_name"))
+
+
+def decode_private_channel(profile: dict, channel_id: str) -> Channel:
+    return Channel(channel_id, ChannelType.DIRECT, profile.get("nickname"))
+
+
+def decode_guild(group: dict) -> Guild:
+    return Guild(str(group["group_id"]), group.get("group_name"), group_avatar(group["group_id"]))
+
+
+def decode_member(member: dict) -> Member:
+    user_id = str(member["user_id"])
+    user = User(user_id, member.get("nickname"), avatar=user_avatar(user_id))
+    joined_at = member.get("join_time")
+    return Member(
+        user=user,
+        nick=member.get("card") or member.get("nickname"),
+        avatar=user_avatar(user_id),
+        joined_at=datetime.fromtimestamp(joined_at) if joined_at else None,
+    )
+
+
+def decode_friend(friend: dict) -> User:
+    user_id = str(friend["user_id"])
+    return User(user_id, friend.get("nickname"), avatar=user_avatar(user_id))
+
+
+def decode_login_user(login: dict) -> User:
+    user_id = str(login["uin"])
+    return User(user_id, login.get("nickname"), avatar=user_avatar(user_id))
+
+
+def decode_user_profile(profile: dict, user_id: str) -> User:
+    return User(user_id, profile.get("nickname"), avatar=user_avatar(user_id), nick=profile.get("remark"))
+
+
+def decode_guild_channel_id(data: dict) -> tuple[str | None, str]:
+    scene = data.get("message_scene")
+    peer_id = str(data.get("peer_id"))
+    if scene == "group":
+        return peer_id, peer_id
+    if scene == "temp":
+        return None, f"private:temp_{peer_id}"
+    return None, f"private:{peer_id}"
+
+
+def get_scene_and_peer(channel_id: str) -> tuple[Literal["friend", "group", "temp"], int]:
+    if channel_id.startswith("private:temp_"):
+        return "temp", int(channel_id.removeprefix("private:temp_"))
+    if channel_id.startswith("private:"):
+        return "friend", int(channel_id.removeprefix("private:"))
+    return "group", int(channel_id)