goyban
a35abb0541
Completes TODO #5. Forwarding a post now asks whether to save it whole or extract its media; sending media directly is unchanged. Built on forwarding rather than copying file_ids. The bot already has the post in its private chat with the owner, so it stores message ids and forwards them back. forwardMessages takes a list and keeps an album intact, which is what made the caption-plus-album case tractable; ids must be strictly ascending, so they are sorted on the way out. - Vault rows live in their own table and every query is scoped to an owner_id, so a guessed id cannot reach someone else's saved post. Nothing in the vault is searchable inline, moderated or federated, and the separate table makes that structural rather than a flag. - Albums arrive as separate messages with no completion signal, so they are buffered by media_group_id behind a 1s debounce and prompt once. - The add path now takes a list of items, so extracting an album stores every piece under one set of keywords instead of asking repeatedly. - /vault searches keywords and the original caption, since the point of a vault is finding something you never got round to tagging. - Retrieval failures are reported: a message id is a pointer into the chat, not a self-contained file_id, so clearing history breaks it. Three near-identical keyword collectors are now one helper. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
1026 lines
38 KiB
Python
1026 lines
38 KiB
Python
import asyncio
|
||
import logging
|
||
import os
|
||
import time
|
||
import uuid
|
||
|
||
from aiogram import Bot, Dispatcher, F
|
||
from aiogram.exceptions import TelegramAPIError
|
||
from aiogram.filters import Command, StateFilter
|
||
from aiogram.fsm.context import FSMContext
|
||
from aiogram.fsm.state import State, StatesGroup
|
||
from aiogram.fsm.storage.memory import MemoryStorage
|
||
from aiogram.types import (
|
||
BotCommand,
|
||
BotCommandScopeChat,
|
||
CallbackQuery,
|
||
InlineKeyboardButton,
|
||
InlineKeyboardMarkup,
|
||
InlineQuery,
|
||
InlineQueryResultCachedAudio,
|
||
InlineQueryResultCachedDocument,
|
||
InlineQueryResultCachedMpeg4Gif,
|
||
InlineQueryResultCachedPhoto,
|
||
InlineQueryResultCachedSticker,
|
||
InlineQueryResultCachedVideo,
|
||
InlineQueryResultCachedVoice,
|
||
InlineQueryResultArticle,
|
||
InlineQueryResultsButton,
|
||
InputTextMessageContent,
|
||
Message,
|
||
)
|
||
|
||
import db
|
||
|
||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||
log = logging.getLogger(__name__)
|
||
|
||
BOT_TOKEN = os.environ["BOT_TOKEN"]
|
||
|
||
ALLOWED_USERS: set[int] = {
|
||
int(uid.strip()) for uid in os.environ.get("ALLOWED_USERS", "").split(",") if uid.strip()
|
||
}
|
||
ADMIN_USERS: set[int] = {
|
||
int(uid.strip()) for uid in os.environ.get("ADMIN_USERS", "").split(",") if uid.strip()
|
||
}
|
||
|
||
def _env_flag(name: str) -> bool:
|
||
return os.environ.get(name, "").strip().lower() in {"1", "true", "yes"}
|
||
|
||
|
||
# Off by default: with ALLOWED_USERS unset, is_allowed() lets everyone in, so
|
||
# opening submissions implicitly would accept media from the whole world.
|
||
ACCEPT_SUBMISSIONS = _env_flag("ACCEPT_SUBMISSIONS")
|
||
|
||
# Public mode swaps the ALLOWED_USERS allowlist for "is this user in the chats
|
||
# I require?". With no REQUIRED_CHATS the bot is simply open to everyone.
|
||
PUBLIC_MODE = _env_flag("PUBLIC_MODE")
|
||
|
||
|
||
def _parse_chat(raw: str) -> int | str:
|
||
"""Telegram takes either @username or a numeric id."""
|
||
raw = raw.strip()
|
||
try:
|
||
return int(raw)
|
||
except ValueError:
|
||
return raw
|
||
|
||
|
||
REQUIRED_CHATS: list[int | str] = [
|
||
_parse_chat(c) for c in os.environ.get("REQUIRED_CHATS", "").split(",") if c.strip()
|
||
]
|
||
# "all" (default): must be in every listed chat. "any": one is enough.
|
||
REQUIRE_MODE = os.environ.get("REQUIRE_MODE", "all").strip().lower()
|
||
|
||
bot = Bot(token=BOT_TOKEN)
|
||
dp = Dispatcher(storage=MemoryStorage())
|
||
|
||
USER_COMMANDS = [
|
||
BotCommand(command="start", description="Start the bot"),
|
||
BotCommand(command="menu", description="Show available commands"),
|
||
BotCommand(command="list", description="Browse saved memes"),
|
||
BotCommand(command="vault", description="Retrieve a saved post"),
|
||
BotCommand(command="cancel", description="Cancel current operation"),
|
||
]
|
||
ADMIN_COMMANDS = [
|
||
*USER_COMMANDS,
|
||
BotCommand(command="review", description="Review pending submissions"),
|
||
BotCommand(command="edit", description="Edit meme keywords"),
|
||
BotCommand(command="delete", description="Remove a meme by ID"),
|
||
]
|
||
|
||
|
||
def is_allowed(user_id: int) -> bool:
|
||
if user_id in ADMIN_USERS:
|
||
return True
|
||
return not ALLOWED_USERS or user_id in ALLOWED_USERS
|
||
|
||
|
||
def is_admin(user_id: int) -> bool:
|
||
if ADMIN_USERS:
|
||
return user_id in ADMIN_USERS
|
||
# An unconfigured private bot treats its users as admins, which is the
|
||
# original single-owner behaviour. That fallback must not survive into
|
||
# public mode, where it would hand admin rights to anyone who found the
|
||
# bot. check_admin_config() flags the missing setting at startup.
|
||
if PUBLIC_MODE:
|
||
return False
|
||
return not ALLOWED_USERS or user_id in ALLOWED_USERS
|
||
|
||
|
||
# Inline queries fire on every keystroke, so an uncached membership check would
|
||
# mean several API calls per typed character. Denials expire quickly so that a
|
||
# user who has just joined is not left waiting for the full window.
|
||
MEMBERSHIP_TTL_ALLOWED = 300.0
|
||
MEMBERSHIP_TTL_DENIED = 30.0
|
||
|
||
ACCEPTED_MEMBER_STATUSES = {"creator", "administrator", "member"}
|
||
|
||
_membership_cache: dict[tuple[int, int | str], tuple[bool, float]] = {}
|
||
|
||
|
||
async def _fetch_membership(user_id: int, chat: int | str) -> bool:
|
||
try:
|
||
member = await bot.get_chat_member(chat, user_id)
|
||
except TelegramAPIError:
|
||
# Usually the bot is not an admin in that chat, or the user is unknown
|
||
# to it. Fail closed; check_required_chats() reports setup problems at
|
||
# startup so this does not fail closed silently forever.
|
||
log.warning("membership check failed for user %s in chat %s", user_id, chat, exc_info=True)
|
||
return False
|
||
|
||
if member.status in ACCEPTED_MEMBER_STATUSES:
|
||
return True
|
||
# A restricted user is still a member unless they have actually left.
|
||
if member.status == "restricted":
|
||
return bool(getattr(member, "is_member", False))
|
||
return False
|
||
|
||
|
||
async def is_member_of(user_id: int, chat: int | str) -> bool:
|
||
key = (user_id, chat)
|
||
cached = _membership_cache.get(key)
|
||
now = time.monotonic()
|
||
if cached is not None and cached[1] > now:
|
||
return cached[0]
|
||
|
||
result = await _fetch_membership(user_id, chat)
|
||
ttl = MEMBERSHIP_TTL_ALLOWED if result else MEMBERSHIP_TTL_DENIED
|
||
_membership_cache[key] = (result, now + ttl)
|
||
return result
|
||
|
||
|
||
async def passes_membership(user_id: int) -> bool:
|
||
if not REQUIRED_CHATS:
|
||
return True
|
||
if REQUIRE_MODE == "any":
|
||
for chat in REQUIRED_CHATS:
|
||
if await is_member_of(user_id, chat):
|
||
return True
|
||
return False
|
||
for chat in REQUIRED_CHATS:
|
||
if not await is_member_of(user_id, chat):
|
||
return False
|
||
return True
|
||
|
||
|
||
async def has_access(user_id: int) -> bool:
|
||
"""Public mode replaces the allowlist with the membership requirement."""
|
||
if is_admin(user_id):
|
||
return True
|
||
if PUBLIC_MODE:
|
||
return await passes_membership(user_id)
|
||
return is_allowed(user_id)
|
||
|
||
|
||
# Resolved once at startup so the join prompt can show real links.
|
||
REQUIRED_CHAT_LINKS: list[tuple[str, str | None]] = []
|
||
|
||
|
||
def check_admin_config() -> None:
|
||
if PUBLIC_MODE and not ADMIN_USERS:
|
||
log.error(
|
||
"PUBLIC_MODE is on but ADMIN_USERS is empty — nobody can moderate, "
|
||
"add or delete. Set ADMIN_USERS to your Telegram user ID."
|
||
)
|
||
|
||
|
||
async def check_required_chats() -> None:
|
||
"""Verify the bot can actually see each required chat, and remember how to
|
||
link to it. A misconfigured chat would otherwise lock out every user with
|
||
nothing in the logs to explain why."""
|
||
REQUIRED_CHAT_LINKS.clear()
|
||
for chat in REQUIRED_CHATS:
|
||
try:
|
||
info = await bot.get_chat(chat)
|
||
except TelegramAPIError:
|
||
log.error(
|
||
"REQUIRED_CHATS: cannot read %s — add the bot to that chat as an "
|
||
"admin, or users will not be able to pass the membership check.",
|
||
chat, exc_info=True,
|
||
)
|
||
REQUIRED_CHAT_LINKS.append((str(chat), None))
|
||
continue
|
||
url = f"https://t.me/{info.username}" if info.username else info.invite_link
|
||
REQUIRED_CHAT_LINKS.append((info.title or str(chat), url))
|
||
if url is None:
|
||
log.warning("REQUIRED_CHATS: no public link for %s; users will only see its name", chat)
|
||
|
||
|
||
def join_prompt() -> str:
|
||
if not REQUIRED_CHAT_LINKS:
|
||
return "You do not have access to this bot."
|
||
joiner = "any one of" if REQUIRE_MODE == "any" else "all of"
|
||
lines = [f"To use this bot you need to join {joiner} these:"]
|
||
for title, url in REQUIRED_CHAT_LINKS:
|
||
lines.append(f"• <a href=\"{url}\">{title}</a>" if url else f"• {title}")
|
||
return "\n".join(lines)
|
||
|
||
|
||
async def register_commands() -> None:
|
||
"""Publish Telegram's command menu with role-appropriate suggestions."""
|
||
if not ADMIN_USERS:
|
||
await bot.set_my_commands(ADMIN_COMMANDS)
|
||
return
|
||
|
||
await bot.set_my_commands(USER_COMMANDS)
|
||
for admin_id in ADMIN_USERS:
|
||
await bot.set_my_commands(
|
||
ADMIN_COMMANDS,
|
||
scope=BotCommandScopeChat(chat_id=admin_id),
|
||
)
|
||
|
||
|
||
class AddMedia(StatesGroup):
|
||
waiting_for_keywords = State()
|
||
|
||
|
||
class EditMedia(StatesGroup):
|
||
waiting_for_new_keywords = State()
|
||
|
||
|
||
class VaultSave(StatesGroup):
|
||
waiting_for_keywords = State()
|
||
|
||
|
||
def menu_text(user_id: int) -> str:
|
||
if is_admin(user_id):
|
||
return (
|
||
"Send me any <b>photo, GIF, video, voice, music, sticker, or file</b>.\n"
|
||
"I'll ask you for keywords next.\n\n"
|
||
"/list — show all saved memes\n"
|
||
"/vault [search] — retrieve a whole post you saved\n"
|
||
"/review — review pending submissions\n"
|
||
"/edit [search] — edit keywords of a meme\n"
|
||
"/delete <id> — remove a meme by ID\n"
|
||
"/cancel — cancel current operation"
|
||
)
|
||
if ACCEPT_SUBMISSIONS:
|
||
return (
|
||
"Use me inline: type <code>@botname keyword</code> in any chat.\n\n"
|
||
"Send me media to submit it — a moderator will review it before it goes live.\n\n"
|
||
"/list — browse all saved memes"
|
||
)
|
||
return (
|
||
"Use me inline: type <code>@botname keyword</code> in any chat.\n\n"
|
||
"/list — browse all saved memes\n"
|
||
"/vault [search] — retrieve a whole post you saved"
|
||
)
|
||
|
||
|
||
async def deny(msg: Message) -> None:
|
||
"""A public bot tells the user how to get in; a private one stays silent,
|
||
as it always has, rather than confirming to strangers that it exists."""
|
||
if PUBLIC_MODE:
|
||
await msg.answer(join_prompt(), parse_mode="HTML")
|
||
|
||
|
||
# ── /start ───────────────────────────────────────────────────────────────────
|
||
|
||
@dp.message(Command("start"))
|
||
async def cmd_start(msg: Message):
|
||
if not await has_access(msg.from_user.id):
|
||
await deny(msg)
|
||
return
|
||
await msg.answer(
|
||
f"{menu_text(msg.from_user.id)}\n/menu — show this list again",
|
||
parse_mode="HTML",
|
||
)
|
||
|
||
|
||
# ── /menu ────────────────────────────────────────────────────────────────────
|
||
|
||
@dp.message(Command("menu"))
|
||
async def cmd_menu(msg: Message):
|
||
if not await has_access(msg.from_user.id):
|
||
await deny(msg)
|
||
return
|
||
await msg.answer(menu_text(msg.from_user.id), parse_mode="HTML")
|
||
|
||
|
||
# ── /cancel ──────────────────────────────────────────────────────────────────
|
||
|
||
@dp.message(Command("cancel"))
|
||
async def cmd_cancel(msg: Message, state: FSMContext):
|
||
if not await has_access(msg.from_user.id):
|
||
return
|
||
await state.clear()
|
||
await msg.answer("Cancelled.")
|
||
|
||
|
||
# ── step 1: receive media (admin only) ───────────────────────────────────────
|
||
|
||
async def start_keyword_prompt(chat_id: int, state: FSMContext, items: list[dict],
|
||
status: str = "approved") -> None:
|
||
"""Begin collecting keywords for one or more media items. Everything in
|
||
`items` ends up sharing the keywords, which is what makes extracting a
|
||
whole album a single interaction."""
|
||
await state.set_state(AddMedia.waiting_for_keywords)
|
||
await state.update_data(items=items, keywords=[], status=status)
|
||
count = f"{len(items)} items" if len(items) > 1 else "it"
|
||
await bot.send_message(
|
||
chat_id,
|
||
f"Got {count}! Now send keywords — one per message or comma-separated.\n"
|
||
"Send /done when finished.",
|
||
parse_mode="HTML",
|
||
)
|
||
|
||
|
||
async def _receive_media(msg: Message, state: FSMContext, file_id: str, file_unique_id: str,
|
||
media_type: str, status: str = "approved"):
|
||
item = {"file_id": file_id, "file_unique_id": file_unique_id, "media_type": media_type}
|
||
await start_keyword_prompt(msg.chat.id, state, [item], status)
|
||
|
||
|
||
# Ordered: Telegram sets `document` alongside `animation` for backward
|
||
# compatibility, so animations must be matched before documents or every GIF
|
||
# would be filed as a file. Round videos are stored as plain "video" so they
|
||
# stay reachable from inline search, which has no cached video_note result.
|
||
MEDIA_EXTRACTORS = (
|
||
("photo", lambda m: m.photo[-1] if m.photo else None),
|
||
("gif", lambda m: m.animation),
|
||
("sticker", lambda m: m.sticker),
|
||
("video", lambda m: m.video),
|
||
("video", lambda m: m.video_note),
|
||
("voice", lambda m: m.voice),
|
||
("audio", lambda m: m.audio),
|
||
("document", lambda m: m.document),
|
||
)
|
||
|
||
MEDIA_FILTER = (
|
||
F.photo | F.animation | F.sticker | F.video
|
||
| F.video_note | F.voice | F.audio | F.document
|
||
)
|
||
|
||
|
||
def extract_media(msg: Message) -> tuple[str, str, str] | None:
|
||
"""Return (media_type, file_id, file_unique_id) for any supported media."""
|
||
for media_type, getter in MEDIA_EXTRACTORS:
|
||
obj = getter(msg)
|
||
if obj is not None:
|
||
return media_type, obj.file_id, obj.file_unique_id
|
||
return None
|
||
|
||
|
||
@dp.message(MEDIA_FILTER, StateFilter(None), F.forward_origin.is_(None))
|
||
async def on_media(msg: Message, state: FSMContext):
|
||
user_id = msg.from_user.id
|
||
if is_admin(user_id):
|
||
status = "approved"
|
||
elif ACCEPT_SUBMISSIONS and await has_access(user_id):
|
||
status = "pending"
|
||
else:
|
||
return
|
||
found = extract_media(msg)
|
||
if found is None:
|
||
return
|
||
media_type, file_id, file_unique_id = found
|
||
await _receive_media(msg, state, file_id, file_unique_id, media_type, status)
|
||
|
||
|
||
# ── step 2: collect keywords ─────────────────────────────────────────────────
|
||
|
||
@dp.message(AddMedia.waiting_for_keywords, Command("done"))
|
||
async def cmd_done(msg: Message, state: FSMContext):
|
||
data = await state.get_data()
|
||
keywords: list[str] = data.get("keywords", [])
|
||
if not keywords:
|
||
await msg.answer("You haven't added any keywords yet. Send some, then /done.")
|
||
return
|
||
status = data.get("status", "approved")
|
||
items = data.get("items", [])
|
||
for item in items:
|
||
await db.add_media(
|
||
msg.from_user.id,
|
||
item["file_id"],
|
||
item["file_unique_id"],
|
||
item["media_type"],
|
||
keywords,
|
||
status=status,
|
||
)
|
||
await state.clear()
|
||
saved = f"{len(items)} items" if len(items) > 1 else "it"
|
||
if status == "pending":
|
||
await msg.answer(
|
||
f"Submitted {saved} with keywords: <b>{', '.join(keywords)}</b>\n"
|
||
"A moderator will review it before it becomes searchable.",
|
||
parse_mode="HTML",
|
||
)
|
||
else:
|
||
await msg.answer(f"Saved {saved} with keywords: <b>{', '.join(keywords)}</b>",
|
||
parse_mode="HTML")
|
||
|
||
|
||
async def accumulate_keywords(msg: Message, state: FSMContext) -> None:
|
||
"""Shared by adding, editing and vault saves: merge in whatever was typed
|
||
and echo the running list back."""
|
||
new_kws = [k.strip().lower() for k in msg.text.split(",") if k.strip()]
|
||
if not new_kws:
|
||
return
|
||
data = await state.get_data()
|
||
existing: list[str] = data.get("keywords", [])
|
||
merged = existing + [k for k in new_kws if k not in existing]
|
||
await state.update_data(keywords=merged)
|
||
await msg.reply(f"Keywords so far: <b>{', '.join(merged)}</b>\nSend more or /done.", parse_mode="HTML")
|
||
|
||
|
||
@dp.message(AddMedia.waiting_for_keywords, F.text)
|
||
async def on_keyword(msg: Message, state: FSMContext):
|
||
await accumulate_keywords(msg, state)
|
||
|
||
|
||
@dp.message(AddMedia.waiting_for_keywords, MEDIA_FILTER)
|
||
async def on_media_during_keywords(msg: Message):
|
||
await msg.reply("Still waiting for keywords. Send /done to save the current item, or /cancel to discard.")
|
||
|
||
|
||
# ── /list ────────────────────────────────────────────────────────────────────
|
||
|
||
LIST_PAGE_SIZE = 30
|
||
|
||
|
||
async def list_page(page: int) -> tuple[str, InlineKeyboardMarkup | None]:
|
||
total = await db.count_media()
|
||
if total == 0:
|
||
return "No memes saved yet.", None
|
||
|
||
total_pages = (total + LIST_PAGE_SIZE - 1) // LIST_PAGE_SIZE
|
||
page = max(0, min(page, total_pages - 1))
|
||
items = await db.list_media(page, LIST_PAGE_SIZE)
|
||
lines = [f"<code>{item['id']}</code> [{item['media_type']}] — {item['keywords']}" for item in items]
|
||
text = f"Page {page + 1} of {total_pages}\n" + "\n".join(lines)
|
||
|
||
buttons = []
|
||
if page > 0:
|
||
buttons.append(InlineKeyboardButton(text="‹ Previous", callback_data=f"list:{page - 1}"))
|
||
if page < total_pages - 1:
|
||
buttons.append(InlineKeyboardButton(text="Next ›", callback_data=f"list:{page + 1}"))
|
||
keyboard = InlineKeyboardMarkup(inline_keyboard=[buttons]) if buttons else None
|
||
return text, keyboard
|
||
|
||
|
||
@dp.message(Command("list"))
|
||
async def cmd_list(msg: Message):
|
||
if not await has_access(msg.from_user.id):
|
||
await deny(msg)
|
||
return
|
||
text, keyboard = await list_page(0)
|
||
await msg.answer(text, parse_mode="HTML", reply_markup=keyboard)
|
||
|
||
|
||
@dp.callback_query(F.data.startswith("list:"))
|
||
async def on_list_page(callback: CallbackQuery):
|
||
if not await has_access(callback.from_user.id):
|
||
await callback.answer("Not allowed.", show_alert=True)
|
||
return
|
||
|
||
try:
|
||
page = int(callback.data.split(":", 1)[1])
|
||
except ValueError:
|
||
await callback.answer("Invalid page.", show_alert=True)
|
||
return
|
||
|
||
text, keyboard = await list_page(page)
|
||
await callback.message.edit_text(text, parse_mode="HTML", reply_markup=keyboard)
|
||
await callback.answer()
|
||
|
||
|
||
# ── /delete ──────────────────────────────────────────────────────────────────
|
||
|
||
@dp.message(Command("delete"))
|
||
async def cmd_delete(msg: Message):
|
||
if not is_admin(msg.from_user.id):
|
||
await msg.answer("Only admins can delete memes.")
|
||
return
|
||
parts = msg.text.split(maxsplit=1)
|
||
if len(parts) < 2:
|
||
await msg.answer("Usage: /delete <id> — use the numeric ID from /list", parse_mode="HTML")
|
||
return
|
||
try:
|
||
media_id = int(parts[1].strip())
|
||
except ValueError:
|
||
await msg.answer("Invalid ID. Use the numeric ID from /list.")
|
||
return
|
||
removed = await db.delete_media(media_id)
|
||
await msg.answer("Deleted." if removed else "Not found.")
|
||
|
||
|
||
# ── /edit ────────────────────────────────────────────────────────────────────
|
||
|
||
@dp.message(Command("edit"))
|
||
async def cmd_edit(msg: Message):
|
||
if not is_admin(msg.from_user.id):
|
||
await msg.answer("Only admins can edit memes.")
|
||
return
|
||
parts = msg.text.split(maxsplit=1)
|
||
query = parts[1].strip() if len(parts) > 1 else ""
|
||
items = await db.search_media(query)
|
||
if not items:
|
||
no_match = f" matching '{query}'" if query else ""
|
||
await msg.answer(f"No memes found{no_match}.")
|
||
return
|
||
buttons = []
|
||
for item in items[:10]:
|
||
label = f"[{item['media_type']}] {item['keywords'][:45]}"
|
||
buttons.append([InlineKeyboardButton(text=label, callback_data=f"edit:{item['id']}")])
|
||
keyboard = InlineKeyboardMarkup(inline_keyboard=buttons)
|
||
await msg.answer("Select a meme to edit its keywords:", reply_markup=keyboard)
|
||
|
||
|
||
@dp.callback_query(F.data.startswith("edit:"))
|
||
async def on_edit_select(callback: CallbackQuery, state: FSMContext):
|
||
if not is_admin(callback.from_user.id):
|
||
await callback.answer("Not allowed.", show_alert=True)
|
||
return
|
||
media_id = int(callback.data.split(":", 1)[1])
|
||
item = await db.get_media_by_id(media_id)
|
||
if not item:
|
||
await callback.answer("Meme not found.", show_alert=True)
|
||
return
|
||
await state.set_state(EditMedia.waiting_for_new_keywords)
|
||
await state.update_data(media_id=media_id, keywords=[])
|
||
await callback.message.edit_reply_markup(reply_markup=None)
|
||
await callback.message.answer(
|
||
f"Editing <b>#{item['id']}</b> [{item['media_type']}]\n"
|
||
f"Current keywords: <b>{item['keywords']}</b>\n\n"
|
||
"Send the new keywords (comma-separated or one per message).\n"
|
||
"Send /done when finished, or /cancel to abort.",
|
||
parse_mode="HTML",
|
||
)
|
||
await callback.answer()
|
||
|
||
|
||
@dp.message(EditMedia.waiting_for_new_keywords, Command("done"))
|
||
async def cmd_edit_done(msg: Message, state: FSMContext):
|
||
data = await state.get_data()
|
||
keywords: list[str] = data.get("keywords", [])
|
||
if not keywords:
|
||
await msg.answer("You haven't entered any keywords yet. Send some, then /done.")
|
||
return
|
||
updated = await db.update_media_keywords(data["media_id"], keywords)
|
||
from_review = data.get("from_review", False)
|
||
await state.clear()
|
||
|
||
if not updated:
|
||
await msg.answer("Failed to update (meme may have been deleted).")
|
||
return
|
||
|
||
if from_review:
|
||
# The item is still pending, so the queue hands the same one back with
|
||
# its new keywords — ready to approve.
|
||
await send_next_pending(msg.chat.id, "Keywords updated.")
|
||
return
|
||
|
||
await msg.answer(f"Updated keywords: <b>{', '.join(keywords)}</b>", parse_mode="HTML")
|
||
|
||
|
||
@dp.message(EditMedia.waiting_for_new_keywords, Command("cancel"))
|
||
async def cmd_edit_cancel(msg: Message, state: FSMContext):
|
||
await state.clear()
|
||
await msg.answer("Edit cancelled.")
|
||
|
||
|
||
@dp.message(EditMedia.waiting_for_new_keywords, F.text)
|
||
async def on_edit_keyword(msg: Message, state: FSMContext):
|
||
await accumulate_keywords(msg, state)
|
||
|
||
|
||
# ── forwards: save whole, or extract the media ───────────────────────────────
|
||
#
|
||
# A forwarded album arrives as several separate messages sharing a
|
||
# media_group_id, with no signal that the last one has landed, so they are
|
||
# buffered briefly and treated as one post.
|
||
|
||
ALBUM_DEBOUNCE = 1.0
|
||
PENDING_FORWARD_LIMIT = 200
|
||
|
||
_album_buffers: dict[str, dict] = {}
|
||
_pending_forwards: dict[str, dict] = {}
|
||
|
||
|
||
def _new_forward_record(msg: Message) -> dict:
|
||
return {
|
||
"user_id": msg.from_user.id,
|
||
"chat_id": msg.chat.id,
|
||
"message_ids": [],
|
||
"caption": "",
|
||
"media": [],
|
||
}
|
||
|
||
|
||
def _absorb(record: dict, msg: Message) -> None:
|
||
record["message_ids"].append(msg.message_id)
|
||
# In an album only one message carries the caption.
|
||
text = msg.caption or msg.text or ""
|
||
if text and not record["caption"]:
|
||
record["caption"] = text
|
||
found = extract_media(msg)
|
||
if found is not None:
|
||
media_type, file_id, file_unique_id = found
|
||
record["media"].append(
|
||
{"media_type": media_type, "file_id": file_id, "file_unique_id": file_unique_id}
|
||
)
|
||
|
||
|
||
async def prompt_forward_choice(record: dict) -> None:
|
||
token = uuid.uuid4().hex[:12]
|
||
# Bounded so a long-running bot cannot accumulate abandoned prompts.
|
||
while len(_pending_forwards) >= PENDING_FORWARD_LIMIT:
|
||
_pending_forwards.pop(next(iter(_pending_forwards)))
|
||
_pending_forwards[token] = record
|
||
|
||
record["message_ids"].sort()
|
||
parts = [f"{len(record['message_ids'])} message(s)"]
|
||
if record["media"]:
|
||
parts.append(f"{len(record['media'])} media")
|
||
summary = ", ".join(parts)
|
||
|
||
row = [InlineKeyboardButton(text="📥 Save whole post", callback_data=f"fwd:save:{token}")]
|
||
if record["media"]:
|
||
row.append(InlineKeyboardButton(text="✂️ Extract media", callback_data=f"fwd:extract:{token}"))
|
||
|
||
await bot.send_message(
|
||
record["chat_id"],
|
||
f"Forwarded post received ({summary}).\n\n"
|
||
"<b>Save whole post</b> keeps it intact in your vault — retrievable only "
|
||
"inside this bot, with /vault.\n"
|
||
"<b>Extract media</b> adds the media to your searchable library as usual.",
|
||
parse_mode="HTML",
|
||
reply_markup=InlineKeyboardMarkup(inline_keyboard=[row]),
|
||
)
|
||
|
||
|
||
async def _flush_album(group_id: str) -> None:
|
||
"""Wait until an album has stopped arriving, then prompt once for all of it."""
|
||
while True:
|
||
record = _album_buffers.get(group_id)
|
||
if record is None:
|
||
return
|
||
remaining = record["deadline"] - time.monotonic()
|
||
if remaining <= 0:
|
||
break
|
||
await asyncio.sleep(remaining)
|
||
|
||
record = _album_buffers.pop(group_id, None)
|
||
if record is not None:
|
||
await prompt_forward_choice(record)
|
||
|
||
|
||
@dp.message(F.forward_origin, StateFilter(None))
|
||
async def on_forwarded(msg: Message, state: FSMContext):
|
||
if not await has_access(msg.from_user.id):
|
||
await deny(msg)
|
||
return
|
||
|
||
group_id = msg.media_group_id
|
||
if group_id is None:
|
||
record = _new_forward_record(msg)
|
||
_absorb(record, msg)
|
||
await prompt_forward_choice(record)
|
||
return
|
||
|
||
record = _album_buffers.get(group_id)
|
||
if record is None:
|
||
record = _new_forward_record(msg)
|
||
_album_buffers[group_id] = record
|
||
record["deadline"] = time.monotonic() + ALBUM_DEBOUNCE
|
||
asyncio.create_task(_flush_album(group_id))
|
||
_absorb(record, msg)
|
||
record["deadline"] = time.monotonic() + ALBUM_DEBOUNCE
|
||
|
||
|
||
@dp.callback_query(F.data.startswith("fwd:"))
|
||
async def on_forward_choice(callback: CallbackQuery, state: FSMContext):
|
||
_, action, token = callback.data.split(":", 2)
|
||
record = _pending_forwards.pop(token, None)
|
||
|
||
if record is None:
|
||
await callback.answer("That prompt expired — forward the post again.", show_alert=True)
|
||
return
|
||
if record["user_id"] != callback.from_user.id:
|
||
await callback.answer("Not yours.", show_alert=True)
|
||
return
|
||
|
||
await callback.message.edit_reply_markup(reply_markup=None)
|
||
|
||
if action == "extract":
|
||
if not await has_access(callback.from_user.id):
|
||
await callback.answer("Not allowed.", show_alert=True)
|
||
return
|
||
status = "approved" if is_admin(callback.from_user.id) else "pending"
|
||
await start_keyword_prompt(record["chat_id"], state, record["media"], status)
|
||
await callback.answer()
|
||
return
|
||
|
||
await state.set_state(VaultSave.waiting_for_keywords)
|
||
await state.update_data(
|
||
chat_id=record["chat_id"],
|
||
message_ids=record["message_ids"],
|
||
caption=record["caption"],
|
||
keywords=[],
|
||
)
|
||
hint = "Send keywords, then /done." if record["caption"] else \
|
||
"Send keywords, then /done. This post has no caption, so keywords are how you will find it."
|
||
await callback.message.answer(hint)
|
||
await callback.answer()
|
||
|
||
|
||
@dp.message(VaultSave.waiting_for_keywords, Command("done"))
|
||
async def cmd_vault_done(msg: Message, state: FSMContext):
|
||
data = await state.get_data()
|
||
keywords: list[str] = data.get("keywords", [])
|
||
caption: str = data.get("caption", "")
|
||
|
||
if not keywords and not caption:
|
||
await msg.answer("Add at least one keyword — there is no caption to find this by.")
|
||
return
|
||
|
||
item_id = await db.add_vault_item(
|
||
msg.from_user.id, data["chat_id"], data["message_ids"], caption, keywords
|
||
)
|
||
await state.clear()
|
||
await msg.answer(
|
||
f"Saved to your vault as <b>#{item_id}</b>.\n\n"
|
||
"Vault posts come back whole, exactly as forwarded — but only here in "
|
||
"this chat, with /vault. Inline search does not reach them.",
|
||
parse_mode="HTML",
|
||
)
|
||
|
||
|
||
@dp.message(VaultSave.waiting_for_keywords, F.text)
|
||
async def on_vault_keyword(msg: Message, state: FSMContext):
|
||
await accumulate_keywords(msg, state)
|
||
|
||
|
||
# ── /vault ───────────────────────────────────────────────────────────────────
|
||
|
||
def vault_label(item: dict) -> str:
|
||
text = item["keywords"].replace(",", ", ") or item["caption"]
|
||
text = " ".join(text.split())
|
||
if len(text) > 40:
|
||
text = text[:39] + "…"
|
||
count = len(db.unpack_message_ids(item["message_ids"]))
|
||
suffix = f" ({count})" if count > 1 else ""
|
||
return f"#{item['id']} {text or 'untitled'}{suffix}"
|
||
|
||
|
||
@dp.message(Command("vault"))
|
||
async def cmd_vault(msg: Message):
|
||
if not await has_access(msg.from_user.id):
|
||
await deny(msg)
|
||
return
|
||
|
||
parts = msg.text.split(maxsplit=1)
|
||
query = parts[1].strip() if len(parts) > 1 else ""
|
||
items = await db.search_vault(msg.from_user.id, query)
|
||
|
||
if not items:
|
||
if query:
|
||
await msg.answer(f"Nothing in your vault matching '{query}'.")
|
||
else:
|
||
await msg.answer(
|
||
"Your vault is empty. Forward a post to me and choose "
|
||
"<b>Save whole post</b> to put it here.",
|
||
parse_mode="HTML",
|
||
)
|
||
return
|
||
|
||
rows = [
|
||
[
|
||
InlineKeyboardButton(text=vault_label(item), callback_data=f"vault:get:{item['id']}"),
|
||
InlineKeyboardButton(text="🗑", callback_data=f"vault:del:{item['id']}"),
|
||
]
|
||
for item in items
|
||
]
|
||
total = await db.count_vault(msg.from_user.id)
|
||
header = f"Showing {len(items)} of {total} saved posts." if not query else \
|
||
f"{len(items)} match '{query}'."
|
||
await msg.answer(header, reply_markup=InlineKeyboardMarkup(inline_keyboard=rows))
|
||
|
||
|
||
@dp.callback_query(F.data.startswith("vault:"))
|
||
async def on_vault_action(callback: CallbackQuery):
|
||
_, action, raw_id = callback.data.split(":", 2)
|
||
# Scoped by owner, so a guessed id cannot reach someone else's vault.
|
||
item = await db.get_vault_item(callback.from_user.id, int(raw_id))
|
||
if item is None:
|
||
await callback.answer("Not found.", show_alert=True)
|
||
return
|
||
|
||
if action == "del":
|
||
await db.delete_vault_item(callback.from_user.id, item["id"])
|
||
await callback.answer("Removed from vault.")
|
||
await callback.message.answer(f"Removed #{item['id']} from your vault.")
|
||
return
|
||
|
||
message_ids = db.unpack_message_ids(item["message_ids"])
|
||
try:
|
||
# forwardMessages keeps an album together and preserves the original
|
||
# "forwarded from" header. Telegram requires ascending ids.
|
||
await bot.forward_messages(
|
||
chat_id=callback.message.chat.id,
|
||
from_chat_id=item["chat_id"],
|
||
message_ids=sorted(message_ids),
|
||
)
|
||
await callback.answer()
|
||
except TelegramAPIError:
|
||
log.exception("could not forward vault item %s", item["id"])
|
||
await callback.answer("Could not retrieve that post.", show_alert=True)
|
||
await callback.message.answer(
|
||
f"#{item['id']} could not be retrieved — the original messages may have "
|
||
"been deleted from this chat. Forward the post again to re-save it."
|
||
)
|
||
|
||
|
||
# ── /review ──────────────────────────────────────────────────────────────────
|
||
|
||
# Sending a stored item back needs one call per media type. Stickers cannot
|
||
# carry a caption, so the item and its review buttons are sent as two messages
|
||
# rather than special-casing captions per type.
|
||
MEDIA_SENDERS = {
|
||
"photo": lambda chat_id, fid: bot.send_photo(chat_id, fid),
|
||
"gif": lambda chat_id, fid: bot.send_animation(chat_id, fid),
|
||
"video": lambda chat_id, fid: bot.send_video(chat_id, fid),
|
||
"voice": lambda chat_id, fid: bot.send_voice(chat_id, fid),
|
||
"audio": lambda chat_id, fid: bot.send_audio(chat_id, fid),
|
||
"sticker": lambda chat_id, fid: bot.send_sticker(chat_id, fid),
|
||
"document": lambda chat_id, fid: bot.send_document(chat_id, fid),
|
||
}
|
||
|
||
|
||
def review_keyboard(media_id: int) -> InlineKeyboardMarkup:
|
||
return InlineKeyboardMarkup(inline_keyboard=[
|
||
[
|
||
InlineKeyboardButton(text="✅ Approve", callback_data=f"review:approve:{media_id}"),
|
||
InlineKeyboardButton(text="❌ Reject", callback_data=f"review:reject:{media_id}"),
|
||
],
|
||
[
|
||
InlineKeyboardButton(text="✏️ Fix keywords", callback_data=f"review:edit:{media_id}"),
|
||
InlineKeyboardButton(text="⏸ Stop", callback_data="review:stop:0"),
|
||
],
|
||
])
|
||
|
||
|
||
async def send_next_pending(chat_id: int, notice: str = "") -> None:
|
||
"""Show the next item awaiting review, or report the queue is empty.
|
||
|
||
Nothing tracks where a moderator got to — each call asks the database for
|
||
the oldest pending item. Stopping and resuming later needs no state, and
|
||
two moderators working at once simply take items off the same queue.
|
||
"""
|
||
item = await db.next_pending()
|
||
prefix = f"{notice}\n\n" if notice else ""
|
||
|
||
if item is None:
|
||
await bot.send_message(chat_id, f"{prefix}Nothing left to review.")
|
||
return
|
||
|
||
remaining = await db.count_pending()
|
||
send = MEDIA_SENDERS.get(item["media_type"])
|
||
if send is not None:
|
||
try:
|
||
await send(chat_id, item["file_id"])
|
||
except Exception:
|
||
# A stale or unusable file_id must not wedge the queue — show the
|
||
# details anyway so the moderator can still reject it.
|
||
log.exception("could not send pending item %s for review", item["id"])
|
||
await bot.send_message(chat_id, "⚠️ Could not load this media.")
|
||
else:
|
||
log.warning("no sender for media_type %r (item %s)", item["media_type"], item["id"])
|
||
await bot.send_message(chat_id, f"⚠️ Unsupported media type: {item['media_type']}")
|
||
|
||
await bot.send_message(
|
||
chat_id,
|
||
f"{prefix}Pending <b>#{item['id']}</b> [{item['media_type']}]\n"
|
||
f"From user <code>{item['owner_id']}</code>\n"
|
||
f"Keywords: <b>{item['keywords'].replace(',', ', ')}</b>\n\n"
|
||
f"{remaining} awaiting review.",
|
||
parse_mode="HTML",
|
||
reply_markup=review_keyboard(item["id"]),
|
||
)
|
||
|
||
|
||
@dp.message(Command("review"))
|
||
async def cmd_review(msg: Message):
|
||
if not is_admin(msg.from_user.id):
|
||
await msg.answer("Only moderators can review submissions.")
|
||
return
|
||
await send_next_pending(msg.chat.id)
|
||
|
||
|
||
@dp.callback_query(F.data.startswith("review:"))
|
||
async def on_review_action(callback: CallbackQuery, state: FSMContext):
|
||
if not is_admin(callback.from_user.id):
|
||
await callback.answer("Not allowed.", show_alert=True)
|
||
return
|
||
|
||
_, action, raw_id = callback.data.split(":", 2)
|
||
await callback.message.edit_reply_markup(reply_markup=None)
|
||
|
||
if action == "stop":
|
||
await callback.message.answer("Review paused. Send /review to pick up where you left off.")
|
||
await callback.answer()
|
||
return
|
||
|
||
media_id = int(raw_id)
|
||
|
||
if action == "edit":
|
||
item = await db.get_media_by_id(media_id)
|
||
if not item:
|
||
await send_next_pending(callback.message.chat.id, "That item is gone.")
|
||
await callback.answer()
|
||
return
|
||
await state.set_state(EditMedia.waiting_for_new_keywords)
|
||
await state.update_data(media_id=media_id, keywords=[], from_review=True)
|
||
await callback.message.answer(
|
||
f"Editing <b>#{item['id']}</b>\n"
|
||
f"Current keywords: <b>{item['keywords']}</b>\n\n"
|
||
"Send the new keywords, then /done. /cancel stops reviewing.",
|
||
parse_mode="HTML",
|
||
)
|
||
await callback.answer()
|
||
return
|
||
|
||
status = "approved" if action == "approve" else "rejected"
|
||
changed = await db.set_status(media_id, status, callback.from_user.id)
|
||
notice = f"#{media_id} {status}." if changed else f"#{media_id} was already handled."
|
||
await send_next_pending(callback.message.chat.id, notice)
|
||
await callback.answer()
|
||
|
||
|
||
# Cached inline results, keyed by the media_type stored in the database.
|
||
# Audio carries its own title metadata; video and document require one.
|
||
INLINE_BUILDERS = {
|
||
"photo": lambda rid, fid, kw: InlineQueryResultCachedPhoto(
|
||
id=rid, photo_file_id=fid, title=kw, description=kw),
|
||
"gif": lambda rid, fid, kw: InlineQueryResultCachedMpeg4Gif(
|
||
id=rid, mpeg4_file_id=fid, title=kw),
|
||
"video": lambda rid, fid, kw: InlineQueryResultCachedVideo(
|
||
id=rid, video_file_id=fid, title=kw, description=kw),
|
||
"voice": lambda rid, fid, kw: InlineQueryResultCachedVoice(
|
||
id=rid, voice_file_id=fid, title=kw),
|
||
"sticker": lambda rid, fid, kw: InlineQueryResultCachedSticker(
|
||
id=rid, sticker_file_id=fid),
|
||
"audio": lambda rid, fid, kw: InlineQueryResultCachedAudio(
|
||
id=rid, audio_file_id=fid),
|
||
"document": lambda rid, fid, kw: InlineQueryResultCachedDocument(
|
||
id=rid, document_file_id=fid, title=kw or "file", description=kw),
|
||
}
|
||
|
||
|
||
# ── inline mode ──────────────────────────────────────────────────────────────
|
||
|
||
@dp.inline_query()
|
||
async def inline_handler(query: InlineQuery):
|
||
if not await has_access(query.from_user.id):
|
||
# A bar above the results beats silently returning nothing: tapping it
|
||
# opens the bot in private chat, where /start explains what to join.
|
||
button = None
|
||
if PUBLIC_MODE and REQUIRED_CHATS:
|
||
button = InlineQueryResultsButton(
|
||
text="Tap to see what you need to join", start_parameter="join"
|
||
)
|
||
await query.answer([], cache_time=1, is_personal=True, button=button)
|
||
return
|
||
|
||
try:
|
||
results_raw = await db.search_media(query.query)
|
||
results = []
|
||
|
||
for item in results_raw:
|
||
rid = str(uuid.uuid4())
|
||
kw_display = item["keywords"].replace(",", ", ")
|
||
mtype = item["media_type"]
|
||
|
||
builder = INLINE_BUILDERS.get(mtype)
|
||
if builder is not None:
|
||
results.append(builder(rid, item["file_id"], kw_display))
|
||
|
||
if not results:
|
||
results = [InlineQueryResultArticle(
|
||
id="noop", title="No results found",
|
||
input_message_content=InputTextMessageContent(message_text="—"),
|
||
description=f"No memes matching '{query.query}'",
|
||
)]
|
||
|
||
await query.answer(results, cache_time=5, is_personal=True)
|
||
|
||
except Exception:
|
||
log.exception("inline_handler error for user %s query %r", query.from_user.id, query.query)
|
||
await query.answer([], cache_time=1)
|
||
|
||
|
||
# ── main ─────────────────────────────────────────────────────────────────────
|
||
|
||
async def main():
|
||
await db.init_db()
|
||
check_admin_config()
|
||
await check_required_chats()
|
||
await register_commands()
|
||
log.info("Bot starting…")
|
||
await dp.start_polling(bot)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|