63 lines
2.2 KiB
Python
63 lines
2.2 KiB
Python
|
import asyncio
|
||
|
|
import hashlib
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import ClassVar
|
||
|
|
|
||
|
|
import aiohttp
|
||
|
|
|
||
|
|
from yuxi.utils.logging_config import logger
|
||
|
|
|
||
|
|
|
||
|
|
class MediaPolicy:
|
||
|
|
CACHE_DIR: ClassVar[Path] = Path("/tmp/yuxi_media_cache")
|
||
|
|
SUPPORTED_IMAGE_FORMATS: ClassVar[frozenset[str]] = frozenset({"png", "jpg", "jpeg", "gif", "webp", "bmp"})
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self._download_locks: dict[str, asyncio.Lock] = {}
|
||
|
|
|
||
|
|
def _get_download_lock(self, url: str) -> asyncio.Lock:
|
||
|
|
if url not in self._download_locks:
|
||
|
|
self._download_locks[url] = asyncio.Lock()
|
||
|
|
return self._download_locks[url]
|
||
|
|
|
||
|
|
def validate_size(self, data: bytes, max_size_mb: int = 100) -> bool:
|
||
|
|
limit = max_size_mb * 1024 * 1024
|
||
|
|
return len(data) <= limit
|
||
|
|
|
||
|
|
def detect_format(self, data: bytes) -> str | None:
|
||
|
|
if data[:4] == b"\x89PNG":
|
||
|
|
return "png"
|
||
|
|
if data[:2] == b"\xff\xd8":
|
||
|
|
return "jpg"
|
||
|
|
if data[:6] in (b"GIF87a", b"GIF89a"):
|
||
|
|
return "gif"
|
||
|
|
if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
|
||
|
|
return "webp"
|
||
|
|
if data[:2] == b"BM":
|
||
|
|
return "bmp"
|
||
|
|
return None
|
||
|
|
|
||
|
|
def is_supported_format(self, fmt: str) -> bool:
|
||
|
|
return fmt in self.SUPPORTED_IMAGE_FORMATS
|
||
|
|
|
||
|
|
async def cache_url_content(self, url: str) -> bytes | None:
|
||
|
|
lock = self._get_download_lock(url)
|
||
|
|
async with lock:
|
||
|
|
cache_key = hashlib.md5(url.encode()).hexdigest()[:12]
|
||
|
|
cache_file = self.CACHE_DIR / cache_key
|
||
|
|
if cache_file.exists():
|
||
|
|
return cache_file.read_bytes()
|
||
|
|
|
||
|
|
self.CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||
|
|
try:
|
||
|
|
async with aiohttp.ClientSession() as session:
|
||
|
|
async with session.get(url, timeout=aiohttp.ClientTimeout(total=30)) as resp:
|
||
|
|
if resp.status != 200:
|
||
|
|
return None
|
||
|
|
data = await resp.read()
|
||
|
|
cache_file.write_bytes(data)
|
||
|
|
return data
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Failed to download media URL: {e}")
|
||
|
|
return None
|