89 lines
3.1 KiB
Python
89 lines
3.1 KiB
Python
from datetime import datetime
|
|
|
|
from pyrogram import Client
|
|
from pyrogram.types import Message
|
|
|
|
from userbot.modules.capture import repository
|
|
from userbot.modules.capture.context import CaptureContext
|
|
from userbot.modules.capture.repository import MessageRow
|
|
from userbot.modules.links import extract_links
|
|
from userbot.modules.links import repository as links_repository
|
|
from userbot.modules.media import capture_media, media_unique_id, self_destruct_ttl
|
|
from userbot.modules.raw_json import raw_json
|
|
from utils.policy.models import CaptureToggles
|
|
|
|
|
|
def sender_id(message: Message) -> int | None:
|
|
if message.from_user is not None:
|
|
return message.from_user.id
|
|
if message.sender_chat is not None:
|
|
return message.sender_chat.id
|
|
return None
|
|
|
|
|
|
def callbacks(message: Message) -> list[tuple[int, str | None, bytes | None]]:
|
|
rows = getattr(message.reply_markup, "inline_keyboard", None)
|
|
if not rows:
|
|
return []
|
|
buttons: list[tuple[int, str | None, bytes | None]] = []
|
|
position = 0
|
|
for row in rows:
|
|
for button in row:
|
|
data = button.callback_data
|
|
if data is not None:
|
|
encoded = data.encode() if isinstance(data, str) else data
|
|
buttons.append((position, button.text, encoded))
|
|
position += 1
|
|
return buttons
|
|
|
|
|
|
def message_row(message: Message, chat_id: int, date: datetime) -> MessageRow:
|
|
return MessageRow(
|
|
chat_id=chat_id,
|
|
message_id=message.id,
|
|
date=date,
|
|
sender_id=sender_id(message),
|
|
text=message.text or message.caption,
|
|
raw=raw_json(message),
|
|
has_media=message.media is not None,
|
|
is_self_destruct=self_destruct_ttl(message) is not None,
|
|
edit_date=message.edit_date,
|
|
media_unique_id=media_unique_id(message),
|
|
)
|
|
|
|
|
|
async def capture_message(
|
|
client: Client, message: Message, ctx: CaptureContext, toggles: CaptureToggles
|
|
) -> None:
|
|
if message.empty or message.chat is None or message.date is None:
|
|
return
|
|
chat_id = message.chat.id or 0
|
|
current = await repository.upsert_message(
|
|
ctx.pool, ctx.account_id, message_row(message, chat_id, message.date)
|
|
)
|
|
if ctx.pending_deletes.take(chat_id, message.id):
|
|
await repository.mark_deleted_late(
|
|
ctx.pool, ctx.account_id, chat_id, message.id, message.date
|
|
)
|
|
from userbot.modules.capture.identity import capture_identity # noqa: PLC0415
|
|
|
|
await capture_identity(message, ctx)
|
|
if not current:
|
|
return
|
|
await capture_media(client, message, ctx, chat_id, message.id, toggles)
|
|
buttons = callbacks(message)
|
|
if buttons:
|
|
await repository.insert_callbacks(
|
|
ctx.pool, ctx.account_id, chat_id, message.id, buttons
|
|
)
|
|
|
|
if message.service is not None:
|
|
from userbot.modules.groups import capture_service # noqa: PLC0415
|
|
|
|
await capture_service(client, message, ctx)
|
|
links = extract_links(message)
|
|
if links:
|
|
await links_repository.insert_links(
|
|
ctx.pool, ctx.account_id, chat_id, message.id, links, raw_json(message)
|
|
)
|