from __future__ import annotations import aiohttp from yuxi.channels.exceptions import DeliveryFailedError from yuxi.utils.logging_config import logger _MAX_FILE_SIZE_MB = 100 FILE_TYPE_IMAGE = "1" FILE_TYPE_VOICE = "2" FILE_TYPE_VIDEO = "3" FILE_TYPE_FILE = "4" def validate_media_size(data: bytes, max_size_mb: int = _MAX_FILE_SIZE_MB, label: str = "media") -> None: max_bytes = max_size_mb * 1024 * 1024 actual_size = len(data) if actual_size > max_bytes: raise DeliveryFailedError(f"{label} size {actual_size / 1024 / 1024:.1f}MB exceeds limit of {max_size_mb}MB") async def upload_media( media_data: bytes, token: str, http_client: aiohttp.ClientSession | None = None, filename: str = "media", file_type: str = FILE_TYPE_FILE, group_openid: str | None = None, sandbox: bool = False, ) -> str: api_base = "https://sandbox.api.sgroup.qq.com" if sandbox else "https://api.sgroup.qq.com" headers = {"Authorization": f"QQBot {token}"} content_type_map = { FILE_TYPE_IMAGE: "image/png", FILE_TYPE_VOICE: "audio/mpeg", FILE_TYPE_VIDEO: "video/mp4", FILE_TYPE_FILE: "application/octet-stream", } mime_type = content_type_map.get(file_type, content_type_map[FILE_TYPE_FILE]) form = aiohttp.FormData() form.add_field("file", media_data, filename=filename, content_type=mime_type) form.add_field("file_type", file_type) url = f"{api_base}/v2/groups/{group_openid}/files" if group_openid else f"{api_base}/v2/users/@me/files" async def _do_upload(client: aiohttp.ClientSession) -> str: async with client.post(url, headers=headers, data=form) as resp: if resp.status != 200: raise DeliveryFailedError(f"Media upload failed: HTTP {resp.status}") result = await resp.json() file_id = result.get("file_uuid", "") or result.get("file_info", "") logger.debug(f"[QQBot] Media uploaded, file_type={file_type}, file_id={file_id}") return file_id if http_client: return await _do_upload(http_client) else: async with aiohttp.ClientSession() as session: return await _do_upload(session) async def upload_image( image_data: bytes, token: str, http_client: aiohttp.ClientSession | None = None, filename: str = "image.png", group_openid: str | None = None, sandbox: bool = False, ) -> str: return await upload_media( image_data, token, http_client=http_client, filename=filename, file_type=FILE_TYPE_IMAGE, group_openid=group_openid, sandbox=sandbox, ) async def download_media( file_id: str, token: str, http_client: aiohttp.ClientSession | None = None, sandbox: bool = False, ) -> bytes: api_base = "https://sandbox.api.sgroup.qq.com" if sandbox else "https://api.sgroup.qq.com" headers = {"Authorization": f"QQBot {token}"} async def _do_download(client: aiohttp.ClientSession) -> bytes: async with client.get( f"{api_base}/v2/users/@me/files/{file_id}", headers=headers, ) as resp: if resp.status != 200: raise DeliveryFailedError(f"Media download failed: HTTP {resp.status}") return await resp.read() if http_client: return await _do_download(http_client) else: async with aiohttp.ClientSession() as session: return await _do_download(session) def build_media_payload( chat_id: str, file_id: str, content: str = "", msg_type: int = 7, ) -> dict: from .constants import DM_CHAT_PREFIX, GROUP_CHAT_PREFIX payload: dict = { "msg_type": msg_type, } if msg_type == 1: payload["image"] = file_id elif msg_type == 4: payload["file"] = file_id else: payload["media"] = {"file_info": file_id} if content: payload["content"] = content[:2000] if chat_id.startswith(GROUP_CHAT_PREFIX): payload["group_openid"] = chat_id.replace(GROUP_CHAT_PREFIX, "") elif not chat_id.startswith(DM_CHAT_PREFIX): payload["channel_id"] = chat_id return payload async def download_image( file_id: str, token: str, http_client: aiohttp.ClientSession | None = None, sandbox: bool = False, ) -> bytes: return await download_media(file_id, token, http_client, sandbox) _CHUNK_SIZE = 5 * 1024 * 1024 _UPLOAD_CACHE: dict[str, str] = {} def _make_cache_key(data: bytes) -> str: import hashlib return hashlib.sha256(data).hexdigest() async def upload_media_cached( media_data: bytes, token: str, http_client: aiohttp.ClientSession | None = None, filename: str = "media", file_type: str = FILE_TYPE_FILE, group_openid: str | None = None, sandbox: bool = False, use_cache: bool = True, ) -> str: if use_cache: cache_key = _make_cache_key(media_data) cached = _UPLOAD_CACHE.get(cache_key) if cached: logger.debug("Media upload: cache hit for %s", filename) return cached file_id = await upload_media( media_data, token, http_client=http_client, filename=filename, file_type=file_type, group_openid=group_openid, sandbox=sandbox, ) if use_cache and file_id: cache_key = _make_cache_key(media_data) _UPLOAD_CACHE[cache_key] = file_id return file_id def clear_upload_cache() -> None: _UPLOAD_CACHE.clear() logger.debug("Media upload cache cleared") async def upload_media_chunked( media_data: bytes, token: str, http_client: aiohttp.ClientSession | None = None, filename: str = "media", file_type: str = FILE_TYPE_FILE, group_openid: str | None = None, sandbox: bool = False, chunk_size: int = _CHUNK_SIZE, ) -> str: if len(media_data) <= chunk_size: return await upload_media( media_data, token, http_client=http_client, filename=filename, file_type=file_type, group_openid=group_openid, sandbox=sandbox, ) import math total_chunks = math.ceil(len(media_data) / chunk_size) api_base = "https://sandbox.api.sgroup.qq.com" if sandbox else "https://api.sgroup.qq.com" headers = {"Authorization": f"QQBot {token}"} content_type_map = { FILE_TYPE_IMAGE: "image/png", FILE_TYPE_VOICE: "audio/mpeg", FILE_TYPE_VIDEO: "video/mp4", FILE_TYPE_FILE: "application/octet-stream", } mime_type = content_type_map.get(file_type, content_type_map[FILE_TYPE_FILE]) async def _do_chunked(client: aiohttp.ClientSession) -> str: init_url = f"{api_base}/v2/users/@me/files/chunked" if group_openid: init_url = f"{api_base}/v2/groups/{group_openid}/files/chunked" init_payload = { "filename": filename, "file_type": int(file_type), "total_size": len(media_data), "chunk_size": chunk_size, "total_chunks": total_chunks, } async with client.post( init_url, headers=headers, json=init_payload, ) as resp: if resp.status not in (200, 201): raise DeliveryFailedError(f"Chunked upload init failed: HTTP {resp.status}") init_data = await resp.json() upload_id = init_data.get("upload_id", "") if not upload_id: raise DeliveryFailedError("Chunked upload: no upload_id returned") for i in range(total_chunks): start = i * chunk_size end = min(start + chunk_size, len(media_data)) chunk = media_data[start:end] chunk_url = f"{api_base}/v2/users/@me/files/chunked/{upload_id}" if group_openid: chunk_url = f"{api_base}/v2/groups/{group_openid}/files/chunked/{upload_id}" form = aiohttp.FormData() form.add_field("chunk", chunk, filename=f"{filename}.chunk{i}", content_type=mime_type) form.add_field("chunk_index", str(i)) async with client.post(chunk_url, headers=headers, data=form) as resp: if resp.status not in (200, 201): raise DeliveryFailedError(f"Chunked upload part {i + 1}/{total_chunks} failed: HTTP {resp.status}") complete_url = f"{api_base}/v2/users/@me/files/chunked/{upload_id}/complete" if group_openid: complete_url = f"{api_base}/v2/groups/{group_openid}/files/chunked/{upload_id}/complete" async with client.post(complete_url, headers=headers) as resp: if resp.status != 200: raise DeliveryFailedError(f"Chunked upload complete failed: HTTP {resp.status}") result = await resp.json() return result.get("file_uuid", "") or result.get("file_info", "") if http_client: return await _do_chunked(http_client) else: async with aiohttp.ClientSession() as session: return await _do_chunked(session) async def upload_media_from_url( url: str, token: str, http_client: aiohttp.ClientSession | None = None, filename: str = "media", file_type: str = FILE_TYPE_FILE, group_openid: str | None = None, sandbox: bool = False, ) -> str: async def _do_url_upload(client: aiohttp.ClientSession) -> str: api_base = "https://sandbox.api.sgroup.qq.com" if sandbox else "https://api.sgroup.qq.com" headers = {"Authorization": f"QQBot {token}"} endpoint = f"{api_base}/v2/users/@me/files/url" if group_openid: endpoint = f"{api_base}/v2/groups/{group_openid}/files/url" payload = { "url": url, "file_type": int(file_type), } async with client.post(endpoint, headers=headers, json=payload) as resp: if resp.status not in (200, 201): raise DeliveryFailedError(f"URL upload failed: HTTP {resp.status}") result = await resp.json() return result.get("file_uuid", "") or result.get("file_info", "") if http_client: return await _do_url_upload(http_client) else: async with aiohttp.ClientSession() as session: return await _do_url_upload(session)