| 1 | """QQ channel implementation using botpy SDK. |
| 2 | |
| 3 | Inbound: |
| 4 | - Parse QQ botpy messages (C2C / Group) |
| 5 | - Download attachments to media dir using chunked streaming write (memory-safe) |
| 6 | - Publish to Nanobot bus via BaseChannel._handle_message() |
| 7 | - Content includes a clear, actionable "Received files:" list with local paths |
| 8 | |
| 9 | Outbound: |
| 10 | - Send attachments (msg.media) first via QQ rich media API (base64 upload + msg_type=7) |
| 11 | - Then send text (plain or markdown) |
| 12 | - msg.media supports local paths, file:// paths, and http(s) URLs |
| 13 | |
| 14 | Notes: |
| 15 | - QQ restricts many audio/video formats. We conservatively classify as image vs file. |
| 16 | - Attachment structures differ across botpy versions; we try multiple field candidates. |
| 17 | """ |
| 18 | |
| 19 | from __future__ import annotations |
| 20 | |
| 21 | import asyncio |
| 22 | import base64 |
| 23 | import mimetypes |
| 24 | import os |
| 25 | import re |
| 26 | import time |
| 27 | from collections import deque |
| 28 | from pathlib import Path |
| 29 | from typing import TYPE_CHECKING, Any, Literal |
| 30 | from urllib.parse import unquote, urlparse |
| 31 | |
| 32 | import aiohttp |
| 33 | from loguru import logger |
| 34 | from pydantic import Field |
| 35 | |
| 36 | from nanobot.bus.events import OutboundMessage |
| 37 | from nanobot.bus.queue import MessageBus |
| 38 | from nanobot.channels.base import BaseChannel |
| 39 | from nanobot.config.schema import Base |
| 40 | from nanobot.security.network import validate_url_target |
| 41 | |
| 42 | try: |
| 43 | from nanobot.config.paths import get_media_dir |
| 44 | except Exception: # pragma: no cover |
| 45 | get_media_dir = None # type: ignore |
| 46 | |
| 47 | try: |
| 48 | import botpy |
| 49 | from botpy.http import Route |
| 50 | |
| 51 | QQ_AVAILABLE = True |
| 52 | except ImportError: # pragma: no cover |
| 53 | QQ_AVAILABLE = False |
| 54 | botpy = None |
| 55 | Route = None |
| 56 | |
| 57 | if TYPE_CHECKING: |
| 58 | from botpy.message import BaseMessage, C2CMessage, GroupMessage |
| 59 | from botpy.types.message import Media |
| 60 | |
| 61 | |
| 62 | # QQ rich media file_type: 1=image, 4=file |
| 63 | # (2=voice, 3=video are restricted; we only use image vs file) |
| 64 | QQ_FILE_TYPE_IMAGE = 1 |
| 65 | QQ_FILE_TYPE_FILE = 4 |
| 66 | |
| 67 | _IMAGE_EXTS = { |
| 68 | ".png", |
| 69 | ".jpg", |
| 70 | ".jpeg", |
| 71 | ".gif", |
| 72 | ".bmp", |
| 73 | ".webp", |
| 74 | ".tif", |
| 75 | ".tiff", |
| 76 | ".ico", |
| 77 | ".svg", |
| 78 | } |
| 79 | |
| 80 | # Replace unsafe characters with "_", keep Chinese and common safe punctuation. |
| 81 | _SAFE_NAME_RE = re.compile(r"[^\w.\-()\[\]（）【】\u4e00-\u9fff]+", re.UNICODE) |
| 82 | |
| 83 | |
| 84 | def _sanitize_filename(name: str) -> str: |
| 85 | """Sanitize filename to avoid traversal and problematic chars.""" |
| 86 | name = (name or "").strip() |
| 87 | name = Path(name).name |
| 88 | name = _SAFE_NAME_RE.sub("_", name).strip("._ ") |
| 89 | return name |
| 90 | |
| 91 | |
| 92 | def _is_image_name(name: str) -> bool: |
| 93 | return Path(name).suffix.lower() in _IMAGE_EXTS |
| 94 | |
| 95 | |
| 96 | def _guess_send_file_type(filename: str) -> int: |
| 97 | """Conservative send type: images -> 1, else -> 4.""" |
| 98 | ext = Path(filename).suffix.lower() |
| 99 | mime, _ = mimetypes.guess_type(filename) |
| 100 | if ext in _IMAGE_EXTS or (mime and mime.startswith("image/")): |
| 101 | return QQ_FILE_TYPE_IMAGE |
| 102 | return QQ_FILE_TYPE_FILE |
| 103 | |
| 104 | |
| 105 | def _make_bot_class(channel: QQChannel) -> type[botpy.Client]: |
| 106 | """Create a botpy Client subclass bound to the given channel.""" |
| 107 | intents = botpy.Intents(public_messages=True, direct_message=True) |
| 108 | |
| 109 | class _Bot(botpy.Client): |
| 110 | def __init__(self): |
| 111 | # Disable botpy's file log — nanobot uses loguru; default "botpy.log" fails on read-only fs |
| 112 | super().__init__(intents=intents, ext_handlers=False) |
| 113 | |
| 114 | async def on_ready(self): |
| 115 | logger.info("QQ bot ready: {}", self.robot.name) |
| 116 | |
| 117 | async def on_c2c_message_create(self, message: C2CMessage): |
| 118 | await channel._on_message(message, is_group=False) |
| 119 | |
| 120 | async def on_group_at_message_create(self, message: GroupMessage): |
| 121 | await channel._on_message(message, is_group=True) |
| 122 | |
| 123 | async def on_direct_message_create(self, message): |
| 124 | await channel._on_message(message, is_group=False) |
| 125 | |
| 126 | return _Bot |
| 127 | |
| 128 | |
| 129 | class QQConfig(Base): |
| 130 | """QQ channel configuration using botpy SDK.""" |
| 131 | |
| 132 | enabled: bool = False |
| 133 | app_id: str = "" |
| 134 | secret: str = "" |
| 135 | allow_from: list[str] = Field(default_factory=list) |
| 136 | msg_format: Literal["plain", "markdown"] = "plain" |
| 137 | ack_message: str = "⏳ Processing..." |
| 138 | |
| 139 | # Optional: directory to save inbound attachments. If empty, use nanobot get_media_dir("qq"). |
| 140 | media_dir: str = "" |
| 141 | |
| 142 | # Download tuning |
| 143 | download_chunk_size: int = 1024 * 256 # 256KB |
| 144 | download_max_bytes: int = 1024 * 1024 * 200 # 200MB safety limit |
| 145 | |
| 146 | |
| 147 | class QQChannel(BaseChannel): |
| 148 | """QQ channel using botpy SDK with WebSocket connection.""" |
| 149 | |
| 150 | name = "qq" |
| 151 | display_name = "QQ" |
| 152 | |
| 153 | @classmethod |
| 154 | def default_config(cls) -> dict[str, Any]: |
| 155 | return QQConfig().model_dump(by_alias=True) |
| 156 | |
| 157 | def __init__(self, config: Any, bus: MessageBus): |
| 158 | if isinstance(config, dict): |
| 159 | config = QQConfig.model_validate(config) |
| 160 | super().__init__(config, bus) |
| 161 | self.config: QQConfig = config |
| 162 | |
| 163 | self._client: botpy.Client | None = None |
| 164 | self._http: aiohttp.ClientSession | None = None |
| 165 | |
| 166 | self._processed_ids: deque[str] = deque(maxlen=1000) |
| 167 | self._msg_seq: int = 1 # used to avoid QQ API dedup |
| 168 | self._chat_type_cache: dict[str, str] = {} |
| 169 | |
| 170 | self._media_root: Path = self._init_media_root() |
| 171 | |
| 172 | # --------------------------- |
| 173 | # Lifecycle |
| 174 | # --------------------------- |
| 175 | |
| 176 | def _init_media_root(self) -> Path: |
| 177 | """Choose a directory for saving inbound attachments.""" |
| 178 | if self.config.media_dir: |
| 179 | root = Path(self.config.media_dir).expanduser() |
| 180 | elif get_media_dir: |
| 181 | try: |
| 182 | root = Path(get_media_dir("qq")) |
| 183 | except Exception: |
| 184 | root = Path.home() / ".nanobot" / "media" / "qq" |
| 185 | else: |
| 186 | root = Path.home() / ".nanobot" / "media" / "qq" |
| 187 | |
| 188 | root.mkdir(parents=True, exist_ok=True) |
| 189 | logger.info("QQ media directory: {}", str(root)) |
| 190 | return root |
| 191 | |
| 192 | async def start(self) -> None: |
| 193 | """Start the QQ bot with auto-reconnect loop.""" |
| 194 | if not QQ_AVAILABLE: |
| 195 | logger.error("QQ SDK not installed. Run: pip install qq-botpy") |
| 196 | return |
| 197 | |
| 198 | if not self.config.app_id or not self.config.secret: |
| 199 | logger.error("QQ app_id and secret not configured") |
| 200 | return |
| 201 | |
| 202 | self._running = True |
| 203 | self._http = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=120)) |
| 204 | |
| 205 | self._client = _make_bot_class(self)() |
| 206 | logger.info("QQ bot started (C2C & Group supported)") |
| 207 | await self._run_bot() |
| 208 | |
| 209 | async def _run_bot(self) -> None: |
| 210 | """Run the bot connection with auto-reconnect.""" |
| 211 | while self._running: |
| 212 | try: |
| 213 | await self._client.start(appid=self.config.app_id, secret=self.config.secret) |
| 214 | except Exception as e: |
| 215 | logger.warning("QQ bot error: {}", e) |
| 216 | if self._running: |
| 217 | logger.info("Reconnecting QQ bot in 5 seconds...") |
| 218 | await asyncio.sleep(5) |
| 219 | |
| 220 | async def stop(self) -> None: |
| 221 | """Stop bot and cleanup resources.""" |
| 222 | self._running = False |
| 223 | if self._client: |
| 224 | try: |
| 225 | await self._client.close() |
| 226 | except Exception: |
| 227 | pass |
| 228 | self._client = None |
| 229 | |
| 230 | if self._http: |
| 231 | try: |
| 232 | await self._http.close() |
| 233 | except Exception: |
| 234 | pass |
| 235 | self._http = None |
| 236 | |
| 237 | logger.info("QQ bot stopped") |
| 238 | |
| 239 | # --------------------------- |
| 240 | # Outbound (send) |
| 241 | # --------------------------- |
| 242 | |
| 243 | async def send(self, msg: OutboundMessage) -> None: |
| 244 | """Send attachments first, then text.""" |
| 245 | try: |
| 246 | if not self._client: |
| 247 | logger.warning("QQ client not initialized") |
| 248 | return |
| 249 | |
| 250 | msg_id = msg.metadata.get("message_id") |
| 251 | chat_type = self._chat_type_cache.get(msg.chat_id, "c2c") |
| 252 | is_group = chat_type == "group" |
| 253 | |
| 254 | # 1) Send media |
| 255 | for media_ref in msg.media or []: |
| 256 | ok = await self._send_media( |
| 257 | chat_id=msg.chat_id, |
| 258 | media_ref=media_ref, |
| 259 | msg_id=msg_id, |
| 260 | is_group=is_group, |
| 261 | ) |
| 262 | if not ok: |
| 263 | filename = ( |
| 264 | os.path.basename(urlparse(media_ref).path) |
| 265 | or os.path.basename(media_ref) |
| 266 | or "file" |
| 267 | ) |
| 268 | await self._send_text_only( |
| 269 | chat_id=msg.chat_id, |
| 270 | is_group=is_group, |
| 271 | msg_id=msg_id, |
| 272 | content=f"[Attachment send failed: {filename}]", |
| 273 | ) |
| 274 | |
| 275 | # 2) Send text |
| 276 | if msg.content and msg.content.strip(): |
| 277 | await self._send_text_only( |
| 278 | chat_id=msg.chat_id, |
| 279 | is_group=is_group, |
| 280 | msg_id=msg_id, |
| 281 | content=msg.content.strip(), |
| 282 | ) |
| 283 | except (aiohttp.ClientError, OSError): |
| 284 | # Network / transport errors — propagate so ChannelManager can retry |
| 285 | raise |
| 286 | except Exception: |
| 287 | logger.exception("Error sending QQ message to chat_id={}", msg.chat_id) |
| 288 | |
| 289 | async def _send_text_only( |
| 290 | self, |
| 291 | chat_id: str, |
| 292 | is_group: bool, |
| 293 | msg_id: str | None, |
| 294 | content: str, |
| 295 | ) -> None: |
| 296 | """Send a plain/markdown text message.""" |
| 297 | if not self._client: |
| 298 | return |
| 299 | |
| 300 | self._msg_seq += 1 |
| 301 | use_markdown = self.config.msg_format == "markdown" |
| 302 | payload: dict[str, Any] = { |
| 303 | "msg_type": 2 if use_markdown else 0, |
| 304 | "msg_id": msg_id, |
| 305 | "msg_seq": self._msg_seq, |
| 306 | } |
| 307 | if use_markdown: |
| 308 | payload["markdown"] = {"content": content} |
| 309 | else: |
| 310 | payload["content"] = content |
| 311 | |
| 312 | if is_group: |
| 313 | await self._client.api.post_group_message(group_openid=chat_id, **payload) |
| 314 | else: |
| 315 | await self._client.api.post_c2c_message(openid=chat_id, **payload) |
| 316 | |
| 317 | async def _send_media( |
| 318 | self, |
| 319 | chat_id: str, |
| 320 | media_ref: str, |
| 321 | msg_id: str | None, |
| 322 | is_group: bool, |
| 323 | ) -> bool: |
| 324 | """Read bytes -> base64 upload -> msg_type=7 send.""" |
| 325 | if not self._client: |
| 326 | return False |
| 327 | |
| 328 | data, filename = await self._read_media_bytes(media_ref) |
| 329 | if not data or not filename: |
| 330 | return False |
| 331 | |
| 332 | try: |
| 333 | file_type = _guess_send_file_type(filename) |
| 334 | file_data_b64 = base64.b64encode(data).decode() |
| 335 | |
| 336 | media_obj = await self._post_base64file( |
| 337 | chat_id=chat_id, |
| 338 | is_group=is_group, |
| 339 | file_type=file_type, |
| 340 | file_data=file_data_b64, |
| 341 | file_name=filename, |
| 342 | srv_send_msg=False, |
| 343 | ) |
| 344 | if not media_obj: |
| 345 | logger.error("QQ media upload failed: empty response") |
| 346 | return False |
| 347 | |
| 348 | self._msg_seq += 1 |
| 349 | if is_group: |
| 350 | await self._client.api.post_group_message( |
| 351 | group_openid=chat_id, |
| 352 | msg_type=7, |
| 353 | msg_id=msg_id, |
| 354 | msg_seq=self._msg_seq, |
| 355 | media=media_obj, |
| 356 | ) |
| 357 | else: |
| 358 | await self._client.api.post_c2c_message( |
| 359 | openid=chat_id, |
| 360 | msg_type=7, |
| 361 | msg_id=msg_id, |
| 362 | msg_seq=self._msg_seq, |
| 363 | media=media_obj, |
| 364 | ) |
| 365 | |
| 366 | logger.info("QQ media sent: {}", filename) |
| 367 | return True |
| 368 | except (aiohttp.ClientError, OSError) as e: |
| 369 | # Network / transport errors — propagate for retry by caller |
| 370 | logger.warning("QQ send media network error filename={} err={}", filename, e) |
| 371 | raise |
| 372 | except Exception as e: |
| 373 | # API-level or other non-network errors — return False so send() can fallback |
| 374 | logger.error("QQ send media failed filename={} err={}", filename, e) |
| 375 | return False |
| 376 | |
| 377 | async def _read_media_bytes(self, media_ref: str) -> tuple[bytes | None, str | None]: |
| 378 | """Read bytes from http(s) or local file path; return (data, filename).""" |
| 379 | media_ref = (media_ref or "").strip() |
| 380 | if not media_ref: |
| 381 | return None, None |
| 382 | |
| 383 | # Local file: plain path or file:// URI |
| 384 | if not media_ref.startswith("http://") and not media_ref.startswith("https://"): |
| 385 | try: |
| 386 | if media_ref.startswith("file://"): |
| 387 | parsed = urlparse(media_ref) |
| 388 | # Windows: path in netloc; Unix: path in path |
| 389 | raw = parsed.path or parsed.netloc |
| 390 | local_path = Path(unquote(raw)) |
| 391 | else: |
| 392 | local_path = Path(os.path.expanduser(media_ref)) |
| 393 | |
| 394 | if not local_path.is_file(): |
| 395 | logger.warning("QQ outbound media file not found: {}", str(local_path)) |
| 396 | return None, None |
| 397 | |
| 398 | data = await asyncio.to_thread(local_path.read_bytes) |
| 399 | return data, local_path.name |
| 400 | except Exception as e: |
| 401 | logger.warning("QQ outbound media read error ref={} err={}", media_ref, e) |
| 402 | return None, None |
| 403 | |
| 404 | # Remote URL |
| 405 | ok, err = validate_url_target(media_ref) |
| 406 | if not ok: |
| 407 | logger.warning("QQ outbound media URL validation failed url={} err={}", media_ref, err) |
| 408 | return None, None |
| 409 | |
| 410 | if not self._http: |
| 411 | self._http = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=120)) |
| 412 | try: |
| 413 | async with self._http.get(media_ref, allow_redirects=True) as resp: |
| 414 | if resp.status >= 400: |
| 415 | logger.warning( |
| 416 | "QQ outbound media download failed status={} url={}", |
| 417 | resp.status, |
| 418 | media_ref, |
| 419 | ) |
| 420 | return None, None |
| 421 | data = await resp.read() |
| 422 | if not data: |
| 423 | return None, None |
| 424 | filename = os.path.basename(urlparse(media_ref).path) or "file.bin" |
| 425 | return data, filename |
| 426 | except Exception as e: |
| 427 | logger.warning("QQ outbound media download error url={} err={}", media_ref, e) |
| 428 | return None, None |
| 429 | |
| 430 | # https://github.com/tencent-connect/botpy/issues/198 |
| 431 | # https://bot.q.qq.com/wiki/develop/api-v2/server-inter/message/send-receive/rich-media.html |
| 432 | async def _post_base64file( |
| 433 | self, |
| 434 | chat_id: str, |
| 435 | is_group: bool, |
| 436 | file_type: int, |
| 437 | file_data: str, |
| 438 | file_name: str | None = None, |
| 439 | srv_send_msg: bool = False, |
| 440 | ) -> Media: |
| 441 | """Upload base64-encoded file and return Media object.""" |
| 442 | if not self._client: |
| 443 | raise RuntimeError("QQ client not initialized") |
| 444 | |
| 445 | if is_group: |
| 446 | endpoint = "/v2/groups/{group_openid}/files" |
| 447 | id_key = "group_openid" |
| 448 | else: |
| 449 | endpoint = "/v2/users/{openid}/files" |
| 450 | id_key = "openid" |
| 451 | |
| 452 | payload: dict[str, Any] = { |
| 453 | id_key: chat_id, |
| 454 | "file_type": file_type, |
| 455 | "file_data": file_data, |
| 456 | "srv_send_msg": srv_send_msg, |
| 457 | } |
| 458 | # Only pass file_name for non-image types (file_type=4). |
| 459 | # Passing file_name for images causes QQ client to render them as |
| 460 | # file attachments instead of inline images. |
| 461 | if file_type != QQ_FILE_TYPE_IMAGE and file_name: |
| 462 | payload["file_name"] = file_name |
| 463 | |
| 464 | route = Route("POST", endpoint, **{id_key: chat_id}) |
| 465 | result = await self._client.api._http.request(route, json=payload) |
| 466 | |
| 467 | # Extract only the file_info field to avoid extra fields (file_uuid, ttl, etc.) |
| 468 | # that may confuse QQ client when sending the media object. |
| 469 | if isinstance(result, dict) and "file_info" in result: |
| 470 | return {"file_info": result["file_info"]} |
| 471 | return result |
| 472 | |
| 473 | # --------------------------- |
| 474 | # Inbound (receive) |
| 475 | # --------------------------- |
| 476 | |
| 477 | async def _on_message(self, data: C2CMessage | GroupMessage, is_group: bool = False) -> None: |
| 478 | """Parse inbound message, download attachments, and publish to the bus.""" |
| 479 | try: |
| 480 | if data.id in self._processed_ids: |
| 481 | return |
| 482 | self._processed_ids.append(data.id) |
| 483 | |
| 484 | if is_group: |
| 485 | chat_id = data.group_openid |
| 486 | user_id = data.author.member_openid |
| 487 | self._chat_type_cache[chat_id] = "group" |
| 488 | else: |
| 489 | chat_id = str( |
| 490 | getattr(data.author, "id", None) |
| 491 | or getattr(data.author, "user_openid", "unknown") |
| 492 | ) |
| 493 | user_id = chat_id |
| 494 | self._chat_type_cache[chat_id] = "c2c" |
| 495 | |
| 496 | content = (data.content or "").strip() |
| 497 | |
| 498 | # the data used by tests don't contain attachments property |
| 499 | # so we use getattr with a default of [] to avoid AttributeError in tests |
| 500 | attachments = getattr(data, "attachments", None) or [] |
| 501 | media_paths, recv_lines, att_meta = await self._handle_attachments(attachments) |
| 502 | |
| 503 | # Compose content that always contains actionable saved paths |
| 504 | if recv_lines: |
| 505 | tag = ( |
| 506 | "[Image]" |
| 507 | if any(_is_image_name(Path(p).name) for p in media_paths) |
| 508 | else "[File]" |
| 509 | ) |
| 510 | file_block = "Received files:\n" + "\n".join(recv_lines) |
| 511 | content = ( |
| 512 | f"{content}\n\n{file_block}".strip() if content else f"{tag}\n{file_block}" |
| 513 | ) |
| 514 | |
| 515 | if not content and not media_paths: |
| 516 | return |
| 517 | |
| 518 | if self.config.ack_message: |
| 519 | try: |
| 520 | await self._send_text_only( |
| 521 | chat_id=chat_id, |
| 522 | is_group=is_group, |
| 523 | msg_id=data.id, |
| 524 | content=self.config.ack_message, |
| 525 | ) |
| 526 | except Exception: |
| 527 | logger.debug("QQ ack message failed for chat_id={}", chat_id) |
| 528 | |
| 529 | await self._handle_message( |
| 530 | sender_id=user_id, |
| 531 | chat_id=chat_id, |
| 532 | content=content, |
| 533 | media=media_paths if media_paths else None, |
| 534 | metadata={ |
| 535 | "message_id": data.id, |
| 536 | "attachments": att_meta, |
| 537 | }, |
| 538 | ) |
| 539 | except Exception: |
| 540 | logger.exception("Error handling QQ inbound message id={}", getattr(data, "id", "?")) |
| 541 | |
| 542 | async def _handle_attachments( |
| 543 | self, |
| 544 | attachments: list[BaseMessage._Attachments], |
| 545 | ) -> tuple[list[str], list[str], list[dict[str, Any]]]: |
| 546 | """Extract, download (chunked), and format attachments for agent consumption.""" |
| 547 | media_paths: list[str] = [] |
| 548 | recv_lines: list[str] = [] |
| 549 | att_meta: list[dict[str, Any]] = [] |
| 550 | |
| 551 | if not attachments: |
| 552 | return media_paths, recv_lines, att_meta |
| 553 | |
| 554 | for att in attachments: |
| 555 | url = getattr(att, "url", None) or "" |
| 556 | filename = getattr(att, "filename", None) or "" |
| 557 | ctype = getattr(att, "content_type", None) or "" |
| 558 | |
| 559 | logger.info("Downloading file from QQ: {}", filename or url) |
| 560 | local_path = await self._download_to_media_dir_chunked(url, filename_hint=filename) |
| 561 | |
| 562 | att_meta.append( |
| 563 | { |
| 564 | "url": url, |
| 565 | "filename": filename, |
| 566 | "content_type": ctype, |
| 567 | "saved_path": local_path, |
| 568 | } |
| 569 | ) |
| 570 | |
| 571 | if local_path: |
| 572 | media_paths.append(local_path) |
| 573 | shown_name = filename or os.path.basename(local_path) |
| 574 | recv_lines.append(f"- {shown_name}\n saved: {local_path}") |
| 575 | else: |
| 576 | shown_name = filename or url |
| 577 | recv_lines.append(f"- {shown_name}\n saved: [download failed]") |
| 578 | |
| 579 | return media_paths, recv_lines, att_meta |
| 580 | |
| 581 | async def _download_to_media_dir_chunked( |
| 582 | self, |
| 583 | url: str, |
| 584 | filename_hint: str = "", |
| 585 | ) -> str | None: |
| 586 | """Download an inbound attachment using streaming chunk write. |
| 587 | |
| 588 | Uses chunked streaming to avoid loading large files into memory. |
| 589 | Enforces a max download size and writes to a .part temp file |
| 590 | that is atomically renamed on success. |
| 591 | """ |
| 592 | # Handle protocol-relative URLs (e.g. "//multimedia.nt.qq.com/...") |
| 593 | if url.startswith("//"): |
| 594 | url = f"https:{url}" |
| 595 | |
| 596 | if not self._http: |
| 597 | self._http = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=120)) |
| 598 | |
| 599 | safe = _sanitize_filename(filename_hint) |
| 600 | ts = int(time.time() * 1000) |
| 601 | tmp_path: Path | None = None |
| 602 | |
| 603 | try: |
| 604 | async with self._http.get( |
| 605 | url, |
| 606 | timeout=aiohttp.ClientTimeout(total=120), |
| 607 | allow_redirects=True, |
| 608 | ) as resp: |
| 609 | if resp.status != 200: |
| 610 | logger.warning("QQ download failed: status={} url={}", resp.status, url) |
| 611 | return None |
| 612 | |
| 613 | ctype = (resp.headers.get("Content-Type") or "").lower() |
| 614 | |
| 615 | # Infer extension: url -> filename_hint -> content-type -> fallback |
| 616 | ext = Path(urlparse(url).path).suffix |
| 617 | if not ext: |
| 618 | ext = Path(filename_hint).suffix |
| 619 | if not ext: |
| 620 | if "png" in ctype: |
| 621 | ext = ".png" |
| 622 | elif "jpeg" in ctype or "jpg" in ctype: |
| 623 | ext = ".jpg" |
| 624 | elif "gif" in ctype: |
| 625 | ext = ".gif" |
| 626 | elif "webp" in ctype: |
| 627 | ext = ".webp" |
| 628 | elif "pdf" in ctype: |
| 629 | ext = ".pdf" |
| 630 | else: |
| 631 | ext = ".bin" |
| 632 | |
| 633 | if safe: |
| 634 | if not Path(safe).suffix: |
| 635 | safe = safe + ext |
| 636 | filename = safe |
| 637 | else: |
| 638 | filename = f"qq_file_{ts}{ext}" |
| 639 | |
| 640 | target = self._media_root / filename |
| 641 | if target.exists(): |
| 642 | target = self._media_root / f"{target.stem}_{ts}{target.suffix}" |
| 643 | |
| 644 | tmp_path = target.with_suffix(target.suffix + ".part") |
| 645 | |
| 646 | # Stream write |
| 647 | downloaded = 0 |
| 648 | chunk_size = max(1024, int(self.config.download_chunk_size or 262144)) |
| 649 | max_bytes = max( |
| 650 | 1024 * 1024, int(self.config.download_max_bytes or (200 * 1024 * 1024)) |
| 651 | ) |
| 652 | |
| 653 | def _open_tmp(): |
| 654 | tmp_path.parent.mkdir(parents=True, exist_ok=True) |
| 655 | return open(tmp_path, "wb") # noqa: SIM115 |
| 656 | |
| 657 | f = await asyncio.to_thread(_open_tmp) |
| 658 | try: |
| 659 | async for chunk in resp.content.iter_chunked(chunk_size): |
| 660 | if not chunk: |
| 661 | continue |
| 662 | downloaded += len(chunk) |
| 663 | if downloaded > max_bytes: |
| 664 | logger.warning( |
| 665 | "QQ download exceeded max_bytes={} url={} -> abort", |
| 666 | max_bytes, |
| 667 | url, |
| 668 | ) |
| 669 | return None |
| 670 | await asyncio.to_thread(f.write, chunk) |
| 671 | finally: |
| 672 | await asyncio.to_thread(f.close) |
| 673 | |
| 674 | # Atomic rename |
| 675 | await asyncio.to_thread(os.replace, tmp_path, target) |
| 676 | tmp_path = None # mark as moved |
| 677 | logger.info("QQ file saved: {}", str(target)) |
| 678 | return str(target) |
| 679 | |
| 680 | except Exception as e: |
| 681 | logger.error("QQ download error: {}", e) |
| 682 | return None |
| 683 | finally: |
| 684 | # Cleanup partial file |
| 685 | if tmp_path is not None: |
| 686 | try: |
| 687 | tmp_path.unlink(missing_ok=True) |
| 688 | except Exception: |
| 689 | pass |
| 690 |