ForcePilot/backend/package/yuxi/channel/extensions/wechat-kf/media.py

import os
import tempfile

import httpx

UPLOAD_URL = "/cgi-bin/media/upload"
DOWNLOAD_URL = "/cgi-bin/media/get"
ALLOWED_TYPES = {"image", "voice", "video", "file"}
TYPE_MAX_SIZES = {
    "image": 2 * 1024 * 1024,
    "voice": 2 * 1024 * 1024,
    "video": 10 * 1024 * 1024,
    "file": 20 * 1024 * 1024,
}


class WeChatKFMedia:
    def __init__(self, gateway):
        self._gateway = gateway

    async def upload(self, file_path: str, media_type: str = "image") -> dict:
        if media_type not in ALLOWED_TYPES:
            return {"errcode": -1, "errmsg": f"不支持的类型: {media_type}"}

        token = await self._gateway.get_access_token()
        with open(file_path, "rb") as f:
            resp = await self._gateway._http.post(
                UPLOAD_URL,
                params={"access_token": token, "type": media_type},
                files={"media": (file_path.split("/")[-1], f)},
            )
        return resp.json()

    async def download(self, media_id: str) -> bytes | None:
        token = await self._gateway.get_access_token()
        resp = await self._gateway._http.get(
            DOWNLOAD_URL,
            params={"access_token": token, "media_id": media_id},
        )
        if resp.status_code == 200:
            return resp.content
        return None

    async def upload_from_url(self, image_url: str, media_type: str = "image") -> dict:
        async with httpx.AsyncClient() as client:
            resp = await client.get(image_url)
            if resp.status_code != 200:
                return {"errcode": -1, "errmsg": "下载图片失败"}

        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
            f.write(resp.content)
            tmp_path = f.name

        try:
            return await self.upload(tmp_path, media_type)
        finally:
            os.unlink(tmp_path)