ForcePilot/backend/package/yuxi/channel/extensions/wechat_ilink/outbound.py

from __future__ import annotations

import base64
import logging
import os
import random
import struct
from typing import Any

import httpx

from yuxi.channel.extensions.wechat_ilink.errors import map_api_error

logger = logging.getLogger(__name__)

BASE_URL = "https://ilinkai.weixin.qq.com"
BASE_INFO = {"channel_version": "2.0.0"}
MAX_UTF8_BYTES = 2048


def split_utf8(text: str, max_bytes: int = MAX_UTF8_BYTES) -> list[str]:
    chunks: list[str] = []
    current = ""
    current_bytes = 0

    for char in text:
        char_bytes = len(char.encode("utf-8"))
        if current_bytes + char_bytes > max_bytes:
            chunks.append(current)
            current = char
            current_bytes = char_bytes
        else:
            current += char
            current_bytes += char_bytes

    if current:
        chunks.append(current)
    return chunks


def parse_target(target_id: str) -> tuple[str, str]:
    parts = target_id.split(":", 2)
    if len(parts) >= 3 and parts[1] == "group":
        return ("group", parts[2])
    return ("private", parts[-1])


class ILinkOutbound:
    def __init__(
        self,
        bot_token: str,
        proxy: str | None = None,
        context_store=None,
        base_url: str = BASE_URL,
    ):
        self._bot_token = bot_token
        self._client: httpx.AsyncClient | None = None
        self._proxy = proxy
        self._context_store = context_store
        self._base_url = base_url

    async def start(self) -> None:
        self._client = httpx.AsyncClient(
            base_url=self._base_url,
            timeout=httpx.Timeout(30.0),
            proxy=self._proxy,
        )

    async def stop(self) -> None:
        if self._client:
            await self._client.aclose()
            self._client = None

    @property
    def client(self) -> httpx.AsyncClient:
        if self._client is None:
            raise RuntimeError("ILinkOutbound 未启动")
        return self._client

    def _headers(self) -> dict[str, str]:
        uin_bytes = struct.pack(">I", random.randint(0, 0xFFFFFFFF))
        x_wechat_uin = base64.b64encode(uin_bytes).decode("ascii")
        return {
            "Authorization": f"Bearer {self._bot_token}",
            "AuthorizationType": "ilink_bot_token",
            "X-WECHAT-UIN": x_wechat_uin,
        }

    def _get_context_token(self, user_id: str) -> str:
        if self._context_store:
            return self._context_store.get(user_id)
        return ""

    def _build_payload(
        self,
        user_id: str,
        item_list: list[dict[str, Any]],
    ) -> dict[str, Any]:
        return {
            "msg": {
                "to_user_id": user_id,
                "from_user_id": "",
                "message_type": 2,
                "message_state": 2,
                "context_token": self._get_context_token(user_id),
                "item_list": item_list,
            },
            "base_info": BASE_INFO,
        }

    async def send_text(
        self,
        target_id: str,
        content: str,
        *,
        reply_to_id: str | None = None,
    ) -> list[dict[str, Any]]:
        chat_type, user_id = parse_target(target_id)
        results: list[dict[str, Any]] = []

        chunks = split_utf8(content)
        for chunk in chunks:
            payload = self._build_payload(
                user_id,
                [{"type": 1, "text_item": {"text": chunk}}],
            )
            if chat_type == "group":
                payload["msg"]["chat_type"] = "group"
            if reply_to_id:
                payload["msg"]["reply_to_id"] = reply_to_id

            resp = await self.client.post(
                "/ilink/bot/sendmessage",
                headers=self._headers(),
                json=payload,
            )
            data = self._check_response(resp)
            results.append(data)

        return results

    async def send_image(
        self,
        target_id: str,
        file_path: str,
        *,
        reply_to_id: str | None = None,
    ) -> dict[str, Any]:
        media_id = await self.upload_media("image", file_path)
        return await self._send_media(target_id, "image", media_id, reply_to_id)

    async def send_file(
        self,
        target_id: str,
        file_path: str,
        *,
        reply_to_id: str | None = None,
    ) -> dict[str, Any]:
        media_id = await self.upload_media("file", file_path)
        return await self._send_media(target_id, "file", media_id, reply_to_id)

    async def upload_media(self, media_type: str, file_path: str) -> str:
        import aiofiles

        async with aiofiles.open(file_path, "rb") as f:
            raw_data = await f.read()

        from yuxi.channel.extensions.wechat_ilink.aes_ecb import aes_ecb_encrypt

        filesize = len(raw_data)
        upload_info = await self._get_upload_url(media_type, filesize)
        aes_key = upload_info.get("aeskey", "")
        cdn_url = upload_info.get("cdn_url", "")
        filekey = upload_info.get("filekey", "")

        if aes_key and cdn_url and filekey:
            key_bytes = aes_key.encode("utf-8")
            encrypted_data = aes_ecb_encrypt(raw_data, key_bytes)
            cdn_client = httpx.AsyncClient(timeout=httpx.Timeout(60.0))
            try:
                cdn_resp = await cdn_client.post(
                    f"{cdn_url}/upload",
                    files={"file": (filekey, encrypted_data)},
                )
                cdn_resp.raise_for_status()
            finally:
                await cdn_client.aclose()
            return upload_info.get("encrypt_query_param", "")

        filename = os.path.basename(file_path)
        resp = await self.client.post(
            "/ilink/bot/media/upload",
            headers=self._headers(),
            files={"media": (filename, raw_data)},
            data={"type": media_type},
        )
        data = self._check_response(resp)
        return data["media_id"]

    async def _get_upload_url(self, media_type: str, filesize: int) -> dict[str, Any]:
        body = {
            "media_type": media_type,
            "filesize": filesize,
            "base_info": BASE_INFO,
        }
        resp = await self.client.post(
            "/ilink/bot/getuploadurl",
            headers=self._headers(),
            json=body,
        )
        return self._check_response(resp)

    async def download_media(self, media_id: str) -> bytes:
        resp = await self.client.get(
            "/ilink/bot/media/download",
            headers=self._headers(),
            params={"media_id": media_id},
        )
        resp.raise_for_status()
        return resp.content

    async def _send_media(
        self,
        target_id: str,
        msg_type: str,
        media_id: str,
        reply_to_id: str | None = None,
    ) -> dict[str, Any]:
        chat_type, user_id = parse_target(target_id)
        type_map = {"image": 3, "file": 6, "voice": 4, "video": 5}
        item_type = type_map.get(msg_type, 6)
        payload = self._build_payload(
            user_id,
            [{"type": item_type, "media_item": {"media_id": media_id}}],
        )
        if chat_type == "group":
            payload["msg"]["chat_type"] = "group"
        if reply_to_id:
            payload["msg"]["reply_to_id"] = reply_to_id

        resp = await self.client.post(
            "/ilink/bot/sendmessage",
            headers=self._headers(),
            json=payload,
        )
        return self._check_response(resp)

    @staticmethod
    def _check_response(resp: httpx.Response) -> dict[str, Any]:
        if resp.is_success:
            data = resp.json()
            errcode = data.get("errcode", 0)
            if errcode != 0:
                error = map_api_error(errcode, data.get("errmsg", ""))
                if error:
                    raise error
            return data
        resp.raise_for_status()
        return {}