122 lines
3.4 KiB
Python
122 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import tempfile
|
|
|
|
import httpx
|
|
|
|
from yuxi.channel.extensions.googlechat.api import (
|
|
download_media_content,
|
|
upload_attachment,
|
|
)
|
|
from yuxi.channel.extensions.googlechat.types import ResolvedGoogleChatAccount
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_DEFAULT_MEDIA_MAX_MB = 20
|
|
|
|
|
|
async def download_attachment(
|
|
account: ResolvedGoogleChatAccount,
|
|
resource_name: str,
|
|
max_mb: int = _DEFAULT_MEDIA_MAX_MB,
|
|
) -> tuple[bytes, str, str]:
|
|
max_bytes = max_mb * 1024 * 1024
|
|
content = await download_media_content(account, resource_name, max_bytes)
|
|
|
|
content_type = "application/octet-stream"
|
|
try:
|
|
import magic
|
|
content_type = magic.Magic(mime=True).from_buffer(content)
|
|
except ImportError:
|
|
if content[:4] == b"\xff\xd8\xff":
|
|
content_type = "image/jpeg"
|
|
elif content[:4] == b"\x89PNG":
|
|
content_type = "image/png"
|
|
elif content[:6] in (b"GIF87a", b"GIF89a"):
|
|
content_type = "image/gif"
|
|
|
|
ext = _mime_to_ext(content_type)
|
|
filename = f"attachment{ext}"
|
|
|
|
return content, filename, content_type
|
|
|
|
|
|
async def upload_media_file(
|
|
account: ResolvedGoogleChatAccount,
|
|
space_name: str,
|
|
file_path: str,
|
|
content_type: str | None = None,
|
|
) -> dict:
|
|
if not content_type:
|
|
try:
|
|
import magic
|
|
content_type = magic.Magic(mime=True).from_file(file_path)
|
|
except ImportError:
|
|
content_type = "application/octet-stream"
|
|
|
|
filename = os.path.basename(file_path)
|
|
with open(file_path, "rb") as f:
|
|
content = f.read()
|
|
|
|
return await upload_attachment(account, space_name, filename, content, content_type)
|
|
|
|
|
|
async def fetch_remote_media(url: str, max_mb: int = _DEFAULT_MEDIA_MAX_MB) -> tuple[bytes, str, str]:
|
|
max_bytes = max_mb * 1024 * 1024
|
|
|
|
async with httpx.AsyncClient(timeout=httpx.Timeout(60.0)) as client:
|
|
resp = await client.get(url)
|
|
resp.raise_for_status()
|
|
content = resp.content
|
|
|
|
if len(content) > max_bytes:
|
|
raise ValueError(f"Media size {len(content)} exceeds limit {max_bytes}")
|
|
|
|
content_type = resp.headers.get("content-type", "application/octet-stream")
|
|
ext = _mime_to_ext(content_type)
|
|
filename = f"downloaded{ext}"
|
|
|
|
return content, filename, content_type
|
|
|
|
|
|
def _mime_to_ext(mime: str) -> str:
|
|
ext_map = {
|
|
"image/jpeg": ".jpg",
|
|
"image/png": ".png",
|
|
"image/gif": ".gif",
|
|
"image/webp": ".webp",
|
|
"image/svg+xml": ".svg",
|
|
"video/mp4": ".mp4",
|
|
"video/webm": ".webm",
|
|
"audio/mpeg": ".mp3",
|
|
"audio/wav": ".wav",
|
|
"application/pdf": ".pdf",
|
|
"text/plain": ".txt",
|
|
"text/html": ".html",
|
|
}
|
|
return ext_map.get(mime, ".bin")
|
|
|
|
|
|
def build_attachment_placeholder(attachments: list[dict]) -> str:
|
|
if not attachments:
|
|
return ""
|
|
|
|
types: dict[str, int] = {}
|
|
for att in attachments:
|
|
ct = att.get("contentType", "unknown")
|
|
main_type = ct.split("/")[0]
|
|
types[main_type] = types.get(main_type, 0) + 1
|
|
|
|
parts = []
|
|
for t, count in types.items():
|
|
parts.append(f"{count} {t}")
|
|
|
|
if len(attachments) == 1:
|
|
ct = attachments[0].get("contentType", "")
|
|
main = ct.split("/")[0]
|
|
return f"<media:{main}>"
|
|
|
|
joined = " + ".join(parts)
|
|
return f"[{joined} attached]" |