128 lines
4.6 KiB
Python
128 lines
4.6 KiB
Python
import base64
|
|
import json
|
|
from pathlib import Path
|
|
from urllib.parse import urlparse
|
|
|
|
import httpx
|
|
|
|
|
|
MAX_IMAGES = 4
|
|
MAX_IMAGE_BYTES = 50_000_000
|
|
MAX_TOTAL_IMAGE_BYTES = 50_000_000
|
|
ALLOWED_IMAGE_HOSTS = {"multimedia.nt.qq.com.cn"}
|
|
|
|
|
|
def extract_message(event: dict) -> tuple[str, list[dict]]:
|
|
message = event.get("message")
|
|
if not isinstance(message, list):
|
|
return str(event.get("raw_message") or ""), []
|
|
|
|
parts: list[str] = []
|
|
media: list[dict] = []
|
|
placeholders = {
|
|
"image": "[图片]",
|
|
"video": "[视频]",
|
|
"record": "[语音]",
|
|
"file": "[文件]",
|
|
"face": "[表情]",
|
|
}
|
|
for segment in message:
|
|
kind = str(segment.get("type") or "")
|
|
data = segment.get("data") or {}
|
|
if kind == "text":
|
|
parts.append(str(data.get("text") or ""))
|
|
elif kind == "at":
|
|
parts.append(f"@{data.get('qq') or data.get('user_id') or ''}")
|
|
elif kind in {"image", "video", "record", "file"}:
|
|
if len(media) < MAX_IMAGES:
|
|
media.append({
|
|
"type": kind,
|
|
"file": data.get("file"),
|
|
"url": data.get("url"),
|
|
})
|
|
parts.append(placeholders[kind])
|
|
elif kind in placeholders:
|
|
parts.append(placeholders[kind])
|
|
return "".join(parts).strip(), media
|
|
|
|
|
|
def _detect_image(data: bytes) -> tuple[str, str] | None:
|
|
if data.startswith(bytes.fromhex("ffd8ff")):
|
|
return "image/jpeg", ".jpg"
|
|
if data.startswith(bytes.fromhex("89504e470d0a1a0a")):
|
|
return "image/png", ".png"
|
|
return None
|
|
|
|
|
|
async def fetch_images(media_json: str, cache_dir: str, message_id: int) -> list[dict]:
|
|
try:
|
|
media = json.loads(media_json or "[]")
|
|
except json.JSONDecodeError:
|
|
return []
|
|
candidates = [item for item in media if item.get("type") == "image"][:MAX_IMAGES]
|
|
if not candidates:
|
|
return []
|
|
|
|
root = Path(cache_dir)
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
results: list[dict] = []
|
|
total_bytes = 0
|
|
async with httpx.AsyncClient(timeout=25, follow_redirects=False, trust_env=False) as client:
|
|
for index, item in enumerate(candidates):
|
|
cached = [root / f"message-{message_id}-{index}.jpg",
|
|
root / f"message-{message_id}-{index}.png"]
|
|
existing = next((path for path in cached if path.is_file()), None)
|
|
if existing:
|
|
data = existing.read_bytes()
|
|
detected = _detect_image(data)
|
|
if detected and 0 < len(data) <= MAX_IMAGE_BYTES:
|
|
if total_bytes + len(data) > MAX_TOTAL_IMAGE_BYTES:
|
|
break
|
|
content_type, _ = detected
|
|
results.append({
|
|
"filename": existing.name,
|
|
"content_type": content_type,
|
|
"data_base64": base64.b64encode(data).decode("ascii"),
|
|
"cache_path": str(existing),
|
|
})
|
|
total_bytes += len(data)
|
|
continue
|
|
url = str(item.get("url") or "")
|
|
parsed = urlparse(url)
|
|
if parsed.scheme != "https" or parsed.hostname not in ALLOWED_IMAGE_HOSTS:
|
|
continue
|
|
data_buffer = bytearray()
|
|
async with client.stream("GET", url) as response:
|
|
response.raise_for_status()
|
|
content_length = int(response.headers.get("content-length") or 0)
|
|
if content_length > MAX_IMAGE_BYTES:
|
|
continue
|
|
async for chunk in response.aiter_bytes():
|
|
data_buffer.extend(chunk)
|
|
if len(data_buffer) > MAX_IMAGE_BYTES:
|
|
break
|
|
data = bytes(data_buffer)
|
|
detected = _detect_image(data)
|
|
if not detected or not data or len(data) > MAX_IMAGE_BYTES:
|
|
continue
|
|
if total_bytes + len(data) > MAX_TOTAL_IMAGE_BYTES:
|
|
break
|
|
content_type, suffix = detected
|
|
path = root / f"message-{message_id}-{index}{suffix}"
|
|
path.write_bytes(data)
|
|
results.append({
|
|
"filename": path.name,
|
|
"content_type": content_type,
|
|
"data_base64": base64.b64encode(data).decode("ascii"),
|
|
"cache_path": str(path),
|
|
})
|
|
total_bytes += len(data)
|
|
return results
|
|
|
|
|
|
def cleanup_images(images: list[dict]) -> None:
|
|
for image in images:
|
|
path = image.get("cache_path")
|
|
if path:
|
|
Path(path).unlink(missing_ok=True)
|