feat: add realtime capturing

backend/src/userbot/modules/media/downloader.py

@@ -0,0 +1,71 @@
+from io import BytesIO
+from typing import Any
+
+from pyrogram import Client
+from pyrogram.types import Message
+
+from userbot.modules.capture import repository
+from userbot.modules.capture.context import CaptureContext
+from utils.policy.models import CaptureToggles
+
+_MEDIA_ATTRS = (
+    "photo",
+    "video",
+    "voice",
+    "video_note",
+    "document",
+    "audio",
+    "animation",
+    "sticker",
+)
+
+
+def media_object(message: Message) -> tuple[str | None, Any]:
+    for attr in _MEDIA_ATTRS:
+        obj = getattr(message, attr, None)
+        if obj is not None:
+            return attr, obj
+    return None, None
+
+
+def self_destruct_ttl(message: Message) -> int | None:
+    _, obj = media_object(message)
+    return getattr(obj, "ttl_seconds", None) if obj is not None else None
+
+
+async def capture_media(  # noqa: PLR0913
+    client: Client,
+    message: Message,
+    ctx: CaptureContext,
+    chat_id: int,
+    message_id: int,
+    toggles: CaptureToggles,
+) -> None:
+    kind, obj = media_object(message)
+    if obj is None:
+        return
+    ttl = getattr(obj, "ttl_seconds", None)
+    want = toggles.self_destruct_media if ttl else toggles.media
+    file_size = getattr(obj, "file_size", None)
+    mime = getattr(obj, "mime_type", None)
+    storage_key: str | None = None
+    downloaded = False
+    if want:
+        buffer = await client.download_media(message, in_memory=True)
+        if isinstance(buffer, BytesIO):
+            data = buffer.getvalue()
+            storage_key = ctx.storage.put(data)
+            file_size = len(data)
+            downloaded = True
+    await repository.insert_media(
+        ctx.pool,
+        ctx.account_id,
+        chat_id,
+        message_id,
+        kind or "unknown",
+        storage_key,
+        file_size,
+        mime,
+        ttl,
+        downloaded=downloaded,
+    )