49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
from pyrogram import Client
|
|
from pyrogram.types import Story
|
|
|
|
from userbot.modules.capture.context import CaptureContext
|
|
from userbot.modules.download import download_bytes
|
|
from userbot.modules.raw_json import raw_json
|
|
from userbot.modules.stories import repository
|
|
|
|
|
|
def story_peer_id(story: Story) -> int:
|
|
if story.chat is not None:
|
|
return story.chat.id or 0
|
|
if story.from_user is not None:
|
|
return story.from_user.id or 0
|
|
return 0
|
|
|
|
|
|
async def save_story(client: Client, capture: CaptureContext, story: Story) -> None:
|
|
peer_id = story_peer_id(story)
|
|
storage_key: str | None = None
|
|
file_size: int | None = None
|
|
downloaded = False
|
|
stored = await repository.is_downloaded(
|
|
capture.pool, capture.account_id, peer_id, story.id
|
|
)
|
|
if not (stored or story.deleted or story.media is None):
|
|
data = await download_bytes(client, story)
|
|
if data is not None:
|
|
storage_key = capture.storage.put(data)
|
|
file_size = len(data)
|
|
downloaded = True
|
|
await repository.upsert_story(
|
|
capture.pool,
|
|
capture.account_id,
|
|
peer_id,
|
|
story.id,
|
|
story.date,
|
|
story.expire_date,
|
|
story.caption,
|
|
story.media.name.lower() if story.media else None,
|
|
storage_key,
|
|
file_size,
|
|
story.views,
|
|
raw_json(story.raw),
|
|
pinned=bool(story.pinned),
|
|
deleted=bool(story.deleted),
|
|
downloaded=downloaded,
|
|
)
|