Kris
87fc26cd15
- 调整依赖导入顺序与清理冗余空行 - 修复setup wizard步骤存储逻辑 - 新增消息编辑/转发标记、卡片控件支持 - 完善目标校验、媒体上传、配对存储功能 - 优化流式响应与错误处理逻辑 - 增强SSRF防护与配置灵活性 - 新增/扩展内置命令与卡片处理能力 - 调整媒体大小限制与类型参数
100 lines
3.5 KiB
Python
100 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
import json
|
|
from typing import TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
from googleapiclient.discovery import Resource
|
|
|
|
from yuxi.channels.exceptions import DeliveryFailedError
|
|
from yuxi.utils.logging_config import logger
|
|
|
|
_MAX_DOWNLOAD_SIZE_BYTES = 100 * 1024 * 1024
|
|
|
|
|
|
async def download_media(chat_service: Resource, file_id: str) -> bytes:
|
|
try:
|
|
file_info = json.loads(file_id)
|
|
except (json.JSONDecodeError, TypeError):
|
|
raise DeliveryFailedError(f"Invalid file_id format: {file_id}")
|
|
|
|
resource_name = file_info.get("name", file_info.get("resource_name", ""))
|
|
if not resource_name:
|
|
raise DeliveryFailedError("file_id missing attachment resource name")
|
|
|
|
try:
|
|
attachment = chat_service.spaces().messages().attachments().get(name=resource_name).execute()
|
|
except Exception as e:
|
|
raise DeliveryFailedError(f"Failed to get attachment metadata: {e}")
|
|
|
|
drive_data = attachment.get("driveDataRef", {})
|
|
if drive_data and drive_data.get("driveFileId"):
|
|
try:
|
|
return await _download_drive_file(chat_service, drive_data["driveFileId"])
|
|
except Exception as e:
|
|
raise DeliveryFailedError(f"Drive file download failed (file_id={drive_data.get('driveFileId')}): {e}")
|
|
|
|
try:
|
|
creds = _get_credentials(chat_service)
|
|
token = await _get_access_token(creds)
|
|
|
|
import httpx
|
|
|
|
url = f"https://chat.googleapis.com/v1/{resource_name}?alt=media"
|
|
async with httpx.AsyncClient(timeout=httpx.Timeout(60)) as client:
|
|
resp = await client.get(
|
|
url,
|
|
headers={"Authorization": f"Bearer {token}"},
|
|
follow_redirects=True,
|
|
)
|
|
if resp.status_code == 200:
|
|
content_length = int(resp.headers.get("content-length", 0))
|
|
if content_length > _MAX_DOWNLOAD_SIZE_BYTES:
|
|
raise DeliveryFailedError(
|
|
f"File too large: {content_length} bytes (max {_MAX_DOWNLOAD_SIZE_BYTES})"
|
|
)
|
|
return resp.content
|
|
raise DeliveryFailedError(f"Media download returned HTTP {resp.status_code}")
|
|
except DeliveryFailedError:
|
|
raise
|
|
except Exception as e:
|
|
logger.warning(f"Direct media download failed, falling back: {e}")
|
|
|
|
raise DeliveryFailedError(f"Unable to download media for attachment: {resource_name}")
|
|
|
|
|
|
async def _download_drive_file(chat_service: Resource, drive_file_id: str) -> bytes:
|
|
from googleapiclient.discovery import build
|
|
from googleapiclient.http import MediaIoBaseDownload
|
|
|
|
creds = _get_credentials(chat_service)
|
|
drive_service = build("drive", "v3", credentials=creds, cache_discovery=False)
|
|
|
|
request = drive_service.files().get_media(fileId=drive_file_id)
|
|
buf = io.BytesIO()
|
|
downloader = MediaIoBaseDownload(buf, request)
|
|
done = False
|
|
while not done:
|
|
_, done = downloader.next_chunk()
|
|
return buf.getvalue()
|
|
|
|
|
|
def _get_credentials(chat_service: Resource):
|
|
http = getattr(chat_service, "_http", None)
|
|
if http is None:
|
|
raise DeliveryFailedError("No HTTP transport on chat_service")
|
|
creds = getattr(http, "credentials", None)
|
|
if creds is None:
|
|
raise DeliveryFailedError("No credentials on chat_service HTTP transport")
|
|
return creds
|
|
|
|
|
|
async def _get_access_token(credentials) -> str:
|
|
from google.auth.transport.requests import Request as GARequest
|
|
|
|
if credentials.valid:
|
|
return credentials.token
|
|
credentials.refresh(GARequest())
|
|
return credentials.token
|