ForcePilot/backend/package/yuxi/channel/extensions/qqbot/media.py

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()