from __future__ import annotations import hashlib import logging import os from pathlib import Path from typing import Any logger = logging.getLogger(__name__) MEDIA_SIZE_LIMITS = { "image": 30 * 1024 * 1024, "voice": 20 * 1024 * 1024, "video": 100 * 1024 * 1024, "file": 100 * 1024 * 1024, } MEDIA_EXTENSIONS = { ".jpg": "image", ".jpeg": "image", ".png": "image", ".gif": "image", ".webp": "image", ".bmp": "image", ".silk": "voice", ".wav": "voice", ".mp3": "voice", ".ogg": "voice", ".aac": "voice", ".flac": "voice", ".mp4": "video", ".mov": "video", ".avi": "video", ".mkv": "video", ".webm": "video", } DEFAULT_MEDIA_ROOT = Path.home() / ".forcepilot" / "qqbot" / "media" class QQBotMedia: def __init__( self, api_client: Any = None, media_root: Path | None = None, url_direct_upload: bool = True, ): self._api_client = api_client self._media_root = media_root or DEFAULT_MEDIA_ROOT self._media_root.mkdir(parents=True, exist_ok=True) self._url_direct_upload = url_direct_upload self._upload_cache: dict[str, str] = {} def detect_media_type(self, path_or_url: str) -> str: ext = os.path.splitext(path_or_url.split("?")[0])[1].lower() if ext in MEDIA_EXTENSIONS: return MEDIA_EXTENSIONS[ext] return "image" def check_size_limit(self, file_path: str | Path, media_type: str) -> bool: path = Path(file_path) if not path.exists(): return False limit = MEDIA_SIZE_LIMITS.get(media_type, 100 * 1024 * 1024) size = path.stat().st_size return size <= limit def is_trusted_path(self, file_path: str | Path) -> bool: path = Path(file_path).resolve() try: path.relative_to(self._media_root.resolve()) return True except ValueError: pass for suffix in path.suffixes: if suffix.lower() in MEDIA_EXTENSIONS: return True return False async def upload_file( self, file_path: str, chat_type: str, target_id: str ) -> dict | None: path = Path(file_path) if not path.exists(): logger.warning("Media file not found: %s", file_path) return None if not self.is_trusted_path(path): logger.warning("Media path not trusted: %s", file_path) return None file_data = path.read_bytes() file_hash = hashlib.sha256(file_data).hexdigest()[:16] if file_hash in self._upload_cache: return {"url": self._upload_cache[file_hash], "cache_hit": True} media_type = self.detect_media_type(file_path) if not self.check_size_limit(file_path, media_type): logger.warning("Media file exceeds size limit: %s", file_path) return None if self._api_client: from yuxi.channel.extensions.qqbot.api_routes import ( FILE_TYPE_FILE, FILE_TYPE_IMAGE, FILE_TYPE_VIDEO, FILE_TYPE_VOICE, ) from yuxi.channel.extensions.qqbot.types import QQBotChatType file_type_map = { "image": FILE_TYPE_IMAGE, "voice": FILE_TYPE_VOICE, "video": FILE_TYPE_VIDEO, "file": FILE_TYPE_FILE, } qt = QQBotChatType.C2C if chat_type in ("c2c", "dm") else QQBotChatType.GROUP ft = file_type_map.get(media_type, FILE_TYPE_FILE) att = await self._api_client.upload_media( target_id, file_data, path.name, ft, qt ) if att and att.url: self._upload_cache[file_hash] = att.url return {"url": att.url, "cache_hit": False} return None def clear_upload_cache(self) -> None: self._upload_cache.clear()