58 lines
1.4 KiB
Python
58 lines
1.4 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
from yuxi.utils.logging_config import logger
|
|
|
|
if TYPE_CHECKING:
|
|
from .client import UrbitClient
|
|
|
|
|
|
async def download_media(
|
|
client: UrbitClient,
|
|
file_url: str,
|
|
timeout: float = 30.0,
|
|
*,
|
|
allow_private: bool = False,
|
|
) -> tuple[bytes, str]:
|
|
from .probe import validate_urbit_url
|
|
|
|
is_valid, warnings = validate_urbit_url(file_url, allow_private=allow_private)
|
|
if not is_valid:
|
|
raise ValueError(f"SSRF blocked: {'; '.join(warnings)}")
|
|
|
|
import httpx
|
|
|
|
try:
|
|
r = await client.http.get(
|
|
file_url,
|
|
timeout=httpx.Timeout(timeout),
|
|
follow_redirects=True,
|
|
)
|
|
r.raise_for_status()
|
|
content_type = r.headers.get("content-type", "application/octet-stream")
|
|
return r.content, content_type
|
|
except Exception as e:
|
|
logger.warning(f"[Urbit] Media download failed from {file_url}: {e}")
|
|
raise
|
|
|
|
|
|
_MEDIA_EXTENSIONS: dict[str, str] = {
|
|
"image/jpeg": ".jpg",
|
|
"image/png": ".png",
|
|
"image/gif": ".gif",
|
|
"image/webp": ".webp",
|
|
"image/svg+xml": ".svg",
|
|
"image/heic": ".heic",
|
|
"image/heif": ".heif",
|
|
"video/mp4": ".mp4",
|
|
"video/webm": ".webm",
|
|
"audio/mpeg": ".mp3",
|
|
"audio/ogg": ".ogg",
|
|
"audio/wav": ".wav",
|
|
}
|
|
|
|
|
|
def get_media_extension(content_type: str) -> str:
|
|
return _MEDIA_EXTENSIONS.get(content_type, ".bin")
|