From a35abb054121c7ff6435c8971ddeb55dee34e090 Mon Sep 17 00:00:00 2001 From: goyban Date: Tue, 1 Sep 2026 22:01:30 +0000 Subject: [PATCH] Add a vault for saving whole forwarded posts 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 --- PLAN.md | 50 +++++--- README.fa.md | 18 +++ README.md | 18 +++ VERSION | 2 +- bot.py | 324 +++++++++++++++++++++++++++++++++++++++++++++----- db.py | 107 ++++++++++++++++- test_db.py | 7 +- test_vault.py | 176 +++++++++++++++++++++++++++ 8 files changed, 650 insertions(+), 52 deletions(-) create mode 100644 test_vault.py diff --git a/PLAN.md b/PLAN.md index 29cdfbb..57b38a2 100644 --- a/PLAN.md +++ b/PLAN.md @@ -7,10 +7,10 @@ Progress on branch `develop` (not pushed): - [x] **#6** full media coverage — `efc2be5` - [x] **#3** moderation queue — `0a8c81d` - [x] **#2** public mode + membership gate — `82b5da4` -- [ ] **#5** save vs. extract (vault) +- [x] **#5** save vs. extract (vault) — forward-based - [ ] **#4** federation -Test suites: `test_db.py`, `test_bot.py`, `test_access.py` — 20 cases, all passing. +Test suites: `test_db.py`, `test_bot.py`, `test_access.py`, `test_vault.py` — 29 cases, all passing. --- @@ -98,29 +98,45 @@ With `PUBLIC_MODE=true`, non-admin submissions land as `status='pending'` from # --- -## #5 — Save vs. extract on forward +## #5 — Save vs. extract on forward — **implemented** -On a forwarded message, ask with inline buttons: **[📥 Save whole post]** / **[✂️ Extract media]**. Extract = today's behavior. Save = the private vault. +Built on forwarding rather than the file_id copy originally planned here, which +removed most of the work. -**Albums are the hard part.** A post with caption + multiple media arrives as *several separate messages* sharing a `media_group_id`, with no "album complete" signal. Standard approach: buffer by `media_group_id` and debounce ~1s after the last message before prompting. Ask once per album, not once per file. +**The bot already holds the post.** When you forward something, Telegram puts a +copy in your private chat with the bot, so the bot can forward *that* back later +by `message_id`. It never needs to reach the original channel, store media, or +re-upload anything. -**Schema (migration v4)** — a separate table, not a `media` column, since one vault item holds N files: +**`forwardMessages` (plural) preserves albums.** It takes up to 100 message ids +and returns them as one intact album, so the caption + N media case that made +this item large is handled by Telegram itself. Ids must be strictly ascending. + +**One table, not two:** ```sql -vault_items(id, owner_id, caption TEXT, keywords TEXT, source_chat TEXT, created_at) -vault_files(item_id, file_id, file_unique_id, media_type, position) +vault_items(id, owner_id, chat_id, message_ids TEXT, caption, keywords, created_at) ``` -Vault rows are **always private** and never touched by the federation sync in #4 — enforce that structurally by keeping them in a different table, so no `visibility` flag can ever be mis-set. +Every vault query takes an `owner_id` and filters on it, so a guessed id cannot +reach another user's post. Vault rows live in their own table and are never +searched by inline mode, never moderated, and never federated. -**Search should cover the caption, not just keywords** — the point of a vault is finding a post you saved without having tagged it. This is where the current `LIKE '%,kw%'` approach (`db.py:43`) stops being adequate; use an FTS5 virtual table over `caption + keywords`. +**Flow.** A forwarded message prompts *Save whole post* / *Extract media* +(extract only appears when there is media). Save collects keywords and stores +the pointer; extract runs the existing add path, now generalised to a list so a +whole album shares one set of keywords. Albums are buffered by `media_group_id` +with a 1s debounce, since Telegram gives no "album complete" signal. -**Open question — recall UX.** Telegram **inline mode cannot send an album**; one inline result sends exactly one message. So a multi-file vault item can't come back through `@bot query` intact. Two options: +**Recall is inside the bot only**, as decided: `/vault [search]` lists matching +posts with a retrieve button and a delete button. Matching covers keywords and +the original caption. The save confirmation tells the user plainly that vault +posts do not appear in inline search. -1. Inline returns the item's *first* media with the caption attached — works everywhere, but loses the other files. -2. A `/vault ` command in the bot's private chat that replies with the full album, which you then forward manually — preserves everything, costs one extra step. - -These aren't exclusive; doing both (1 for speed, 2 for fidelity) is probably right, but it's your call. Single-file vault items are unaffected either way. +**The tradeoff accepted:** a `message_id` is a pointer into the bot's chat with +the owner, not self-contained like a `file_id`. If that history is cleared, the +pointer breaks; retrieval catches the failure and tells the user to forward the +post again. --- @@ -164,14 +180,14 @@ The hub is a separate deployable (small FastAPI + SQLite would do). Do you want | 6 | Full media coverage | 0 (only for the refactor) | S | | 3 | Moderation queue | 0 | M | | 2 | Public mode + membership gate | 0, 3 | M | -| 5 | Save vs. extract + vault | 0 | L | +| 5 | Save vs. extract + vault | 0 | M (was L) | | 4 | Federation | 0, 3, hub built | XL | #6 is shippable on its own and touches nothing else — good first commit on `develop`. #4 should stay last: it's the only item that's hard to reverse once other people are running instances against your hub. ## Decisions — settled 2026-09-01 -1. **#5 recall UX** — inline cannot send an album. For multi-file vault items, the inline result sends the **link to the original post** where one exists; otherwise it tells the user to open the bot and pull the full album there. Single-file items send directly as today. +1. **#5 recall UX** — settled further during implementation: the vault is **forward-based and retrievable only inside the bot**. Inline search does not surface vault posts at all; the save confirmation says so. 2. **#4 hub location** — **separate repository.** Channel name and hub hosting are the owner's call, not decided here. 3. **#6 video_note** — **store as `video`**, accepting the loss of round presentation, so round videos stay inline-searchable. See the risk note below. 4. **#4 architecture** — **shared federation channel accepted.** Media rides the channel, metadata rides the hub API. diff --git a/README.fa.md b/README.fa.md index 3a22f93..8869a9b 100644 --- a/README.fa.md +++ b/README.fa.md @@ -59,10 +59,28 @@ docker compose up -d --build | دستور | توضیح | |---|---| | `/list` | نمایش همه رسانه‌های ذخیره‌شده | +| `/vault [search]` | بازیابی پستی که کامل ذخیره کرده‌اید | | `/review` | بررسی یک‌به‌یک ارسال‌های در انتظار تأیید (فقط ادمین) | | `/delete ` | حذف یک مورد (ID را از `/list` بگیرید) | | `/cancel` | لغو عملیات جاری | +### گاوصندوق — ذخیره کامل یک پست + +اگر رسانه‌ای را مستقیم برای ربات بفرستید، مثل همیشه رفتار می‌کند: رسانه را جدا می‌کند و کلیدواژه می‌خواهد. + +اما اگر پستی را **فوروارد** کنید، ربات می‌پرسد چه می‌خواهید: + +- **ذخیره کل پست** — پست را دست‌نخورده با کپشن در گاوصندوق نگه می‌دارد. آلبوم هم آلبوم می‌ماند. +- **جدا کردن رسانه** — همان رفتار قبلی؛ رسانه به کتابخانه قابل جستجو اضافه می‌شود. + +با `/vault [search]` پست‌های ذخیره‌شده را پیدا کنید؛ جستجو هم روی کلیدواژه‌ها و هم روی متن کپشن اصلی انجام می‌شود. کنار هر نتیجه یک دکمه برای فرستادن دوباره پست و یک دکمه برای حذف آن هست. + +> **پست‌های گاوصندوق فقط داخل خود ربات برمی‌گردند.** حالت inline تلگرام نمی‌تواند پیام موجود را فوروارد کند، پس `@yourbotusername keyword` به گاوصندوق دسترسی ندارد. پست را داخل ربات بگیرید و بعد هرجا خواستید فوروارد کنید. + +گاوصندوق برای هر کاربر جداگانه و خصوصی است: در جدولی جدا از کتابخانه میم‌ها ذخیره می‌شود، هر جستجو محدود به صاحب آن است، و هیچ‌چیزِ آن بررسی یا به اشتراک گذاشته نمی‌شود. + +ذخیره‌سازی به‌جای کپی رسانه، اشاره‌ای به پیام موجود در چت شما با ربات را نگه می‌دارد. اگر تاریخچه چتتان با ربات را پاک کنید، آن پست‌ها دیگر قابل بازیابی نیستند — ربات در این حالت به شما اطلاع می‌دهد تا دوباره فوروارد کنید. + ### عمومی‌کردن ربات با تنظیم `PUBLIC_MODE=true` ربات به‌جای فهرست `ALLOWED_USERS`، به هر کسی که شرط عضویت را داشته باشد اجازه دسترسی می‌دهد: diff --git a/README.md b/README.md index e1b2047..9b76782 100644 --- a/README.md +++ b/README.md @@ -65,11 +65,29 @@ Telegram suggests these commands as you type `/`, with a short description besid | `/start` | Start the bot and show usage instructions | | `/menu` | Show the available commands | | `/list` | Show all saved media | +| `/vault [search]` | Retrieve a whole post you saved | | `/review` | Review pending submissions one by one (admin only) | | `/edit [search]` | Select a meme and edit its keywords (admin only) | | `/delete ` | Remove an entry (get ID from `/list`; admin only) | | `/cancel` | Cancel current add operation | +### The vault — saving whole posts + +Send the bot media directly and it behaves as it always has: it extracts the media and asks for keywords. + +**Forward** a post to it instead and it asks what you want: + +- **Save whole post** — keeps the post intact in your vault, caption and all. An album stays an album. +- **Extract media** — the existing behaviour, adding the media to your searchable library. + +Retrieve saved posts with `/vault [search]`, which matches your keywords and the original caption. Each result has a button to forward it back to you, and one to remove it. + +> **Vault posts only come back inside this bot.** Telegram's inline mode cannot forward an existing message, so `@yourbotusername keyword` does not reach the vault. Retrieve the post in the bot chat, then forward it wherever you want. + +The vault is per-user and private: it is stored separately from the meme library, every lookup is scoped to its owner, and nothing in it is ever moderated or shared. + +Saving stores a pointer to the message in your chat with the bot rather than a copy of the media. If you clear your chat history with the bot, saved posts can no longer be retrieved — the bot tells you when that happens so you can forward them again. + ### Making the bot public Set `PUBLIC_MODE=true` to open the bot to anyone who meets a membership requirement, instead of the `ALLOWED_USERS` list: diff --git a/VERSION b/VERSION index f0bb29e..88c5fb8 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -1.3.0 +1.4.0 diff --git a/bot.py b/bot.py index 253dcc2..47c2080 100644 --- a/bot.py +++ b/bot.py @@ -79,6 +79,7 @@ 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 = [ @@ -238,12 +239,17 @@ 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 photo, GIF, video, voice, music, sticker, or file.\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" @@ -257,7 +263,8 @@ def menu_text(user_id: int) -> str: ) return ( "Use me inline: type @botname keyword in any chat.\n\n" - "/list — browse all saved memes" + "/list — browse all saved memes\n" + "/vault [search] — retrieve a whole post you saved" ) @@ -303,23 +310,28 @@ async def cmd_cancel(msg: Message, state: FSMContext): # ── step 1: receive media (admin only) ─────────────────────────────────────── -async def _receive_media(msg: Message, state: FSMContext, file_id: str, file_unique_id: str, - media_type: str, status: str = "approved"): +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( - file_id=file_id, - file_unique_id=file_unique_id, - media_type=media_type, - keywords=[], - status=status, - ) - await msg.reply( - "Got it! Now send keywords — one per message or comma-separated.\n" + 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 @@ -350,7 +362,7 @@ def extract_media(msg: Message) -> tuple[str, str, str] | None: return None -@dp.message(MEDIA_FILTER, StateFilter(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): @@ -376,27 +388,32 @@ async def cmd_done(msg: Message, state: FSMContext): await msg.answer("You haven't added any keywords yet. Send some, then /done.") return status = data.get("status", "approved") - await db.add_media( - msg.from_user.id, - data["file_id"], - data["file_unique_id"], - data["media_type"], - keywords, - status=status, - ) + 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 with keywords: {', '.join(keywords)}\n" + f"Submitted {saved} with keywords: {', '.join(keywords)}\n" "A moderator will review it before it becomes searchable.", parse_mode="HTML", ) else: - await msg.answer(f"Saved with keywords: {', '.join(keywords)}", parse_mode="HTML") + await msg.answer(f"Saved {saved} with keywords: {', '.join(keywords)}", + parse_mode="HTML") -@dp.message(AddMedia.waiting_for_keywords, F.text) -async def on_keyword(msg: Message, state: FSMContext): +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 @@ -407,6 +424,11 @@ async def on_keyword(msg: Message, state: FSMContext): await msg.reply(f"Keywords so far: {', '.join(merged)}\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.") @@ -560,14 +582,254 @@ async def cmd_edit_cancel(msg: Message, state: FSMContext): @dp.message(EditMedia.waiting_for_new_keywords, F.text) async def on_edit_keyword(msg: Message, state: FSMContext): - new_kws = [k.strip().lower() for k in msg.text.split(",") if k.strip()] - if not new_kws: + 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" + "Save whole post keeps it intact in your vault — retrievable only " + "inside this bot, with /vault.\n" + "Extract media 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() - 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: {', '.join(merged)}\nSend more or /done.", parse_mode="HTML") + 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 #{item_id}.\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 " + "Save whole post 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 ────────────────────────────────────────────────────────────────── diff --git a/db.py b/db.py index c30391c..1f985cd 100644 --- a/db.py +++ b/db.py @@ -4,7 +4,7 @@ import os DB_PATH = os.getenv("DB_PATH", "/data/media.db") # Bump this and add a matching _migrate_to_vN() when the schema changes. -SCHEMA_VERSION = 3 +SCHEMA_VERSION = 4 # ── schema / migrations ────────────────────────────────────────────────────── @@ -69,9 +69,28 @@ async def _migrate_to_v3(db) -> None: await db.execute("CREATE INDEX IF NOT EXISTS idx_status ON media(status)") +async def _migrate_to_v4(db) -> None: + """The vault: posts kept whole, as pointers to the messages the bot already + holds in its private chat with the owner. Nothing here is ever shared, so it + lives in its own table rather than as a flag on media.""" + await db.execute(""" + CREATE TABLE IF NOT EXISTS vault_items ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + owner_id INTEGER NOT NULL, + chat_id INTEGER NOT NULL, + message_ids TEXT NOT NULL, + caption TEXT NOT NULL DEFAULT '', + keywords TEXT NOT NULL DEFAULT '', + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP + ) + """) + await db.execute("CREATE INDEX IF NOT EXISTS idx_vault_owner ON vault_items(owner_id)") + + MIGRATIONS = { 2: _migrate_to_v2, 3: _migrate_to_v3, + 4: _migrate_to_v4, } @@ -246,3 +265,89 @@ async def set_status(media_id: int, status: str, reviewer_id: int) -> bool: ) await db.commit() return cur.rowcount > 0 + + +# ── vault ──────────────────────────────────────────────────────────────────── +# +# Every query here takes an owner_id and filters on it. A vault item belongs to +# one person and is never searchable, never moderated and never federated; the +# separate table is what makes that structural rather than a flag to remember. + +VAULT_COLUMNS = "id, owner_id, chat_id, message_ids, caption, keywords, created_at" + + +def pack_message_ids(message_ids: list[int]) -> str: + return ",".join(str(m) for m in message_ids) + + +def unpack_message_ids(packed: str) -> list[int]: + return [int(m) for m in packed.split(",") if m.strip()] + + +async def add_vault_item(owner_id: int, chat_id: int, message_ids: list[int], + caption: str, keywords: list[str]) -> int: + kw = ",".join(k.strip().lower() for k in keywords if k.strip()) + async with aiosqlite.connect(DB_PATH) as db: + cur = await db.execute( + """ + INSERT INTO vault_items (owner_id, chat_id, message_ids, caption, keywords) + VALUES (?,?,?,?,?) + """, + (owner_id, chat_id, pack_message_ids(message_ids), caption or "", kw), + ) + await db.commit() + return cur.lastrowid + + +async def search_vault(owner_id: int, query: str, limit: int = 20) -> list[dict]: + """Match keywords by prefix or the original caption anywhere in its text — + the point of a vault is finding a post you never got round to tagging.""" + query = query.strip().lower() + params: list = [owner_id] + clause = "" + if query: + clause = "AND ((',' || keywords || ',') LIKE ? OR LOWER(caption) LIKE ?)" + params += [f"%,{query}%", f"%{query}%"] + params.append(limit) + + async with aiosqlite.connect(DB_PATH) as db: + db.row_factory = aiosqlite.Row + cursor = await db.execute( + f""" + SELECT {VAULT_COLUMNS} FROM vault_items + WHERE owner_id = ? {clause} + ORDER BY created_at DESC, id DESC + LIMIT ? + """, + params, + ) + return [dict(r) for r in await cursor.fetchall()] + + +async def get_vault_item(owner_id: int, item_id: int) -> dict | None: + """Scoped by owner so one user cannot fetch another user's saved post.""" + async with aiosqlite.connect(DB_PATH) as db: + db.row_factory = aiosqlite.Row + cursor = await db.execute( + f"SELECT {VAULT_COLUMNS} FROM vault_items WHERE id = ? AND owner_id = ?", + (item_id, owner_id), + ) + row = await cursor.fetchone() + return dict(row) if row else None + + +async def delete_vault_item(owner_id: int, item_id: int) -> bool: + async with aiosqlite.connect(DB_PATH) as db: + cur = await db.execute( + "DELETE FROM vault_items WHERE id = ? AND owner_id = ?", (item_id, owner_id) + ) + await db.commit() + return cur.rowcount > 0 + + +async def count_vault(owner_id: int) -> int: + async with aiosqlite.connect(DB_PATH) as db: + cursor = await db.execute( + "SELECT COUNT(*) FROM vault_items WHERE owner_id = ?", (owner_id,) + ) + return (await cursor.fetchone())[0] diff --git a/test_db.py b/test_db.py index 3fd38b8..b25e4df 100644 --- a/test_db.py +++ b/test_db.py @@ -47,12 +47,15 @@ async def test_legacy_upgrade(tmp: str): db = load_db(path) await db.init_db() - assert user_version(path) == 3 + # Pinned to the module so adding a migration does not require editing this. + assert user_version(path) == db.SCHEMA_VERSION rows = sorted((r["id"], r["keywords"]) for r in await db.search_media("")) assert rows == [(7, "cat,funny"), (9, "dog")], f"rows not preserved: {rows}" with sqlite3.connect(path) as c: statuses = {r[0] for r in c.execute("SELECT status FROM media")} + tables = {r[0] for r in c.execute("SELECT name FROM sqlite_master WHERE type='table'")} assert statuses == {"approved"}, f"upgraded rows must stay live: {statuses}" + assert "vault_items" in tables, "upgrading should create the vault table too" print(" legacy v1 database upgrades with ids, keywords and visibility intact") @@ -91,7 +94,7 @@ async def test_fresh_and_idempotent(tmp: str): with sqlite3.connect(path) as c: sql = c.execute("SELECT sql FROM sqlite_master WHERE name='media'").fetchone()[0] assert "UNIQUE(owner_id, file_unique_id)" in sql - assert user_version(path) == 3 + assert user_version(path) == db.SCHEMA_VERSION await db.init_db() await db.init_db() diff --git a/test_vault.py b/test_vault.py new file mode 100644 index 0000000..82188dd --- /dev/null +++ b/test_vault.py @@ -0,0 +1,176 @@ +"""Vault tests. Run with: python test_vault.py + +The vault stores pointers to messages the bot already holds, so the things +worth pinning down are that albums are grouped into one post, that retrieval +asks Telegram for the right message ids, and that one user can never reach +another user's saved posts. +""" +import asyncio +import datetime +import importlib +import logging +import os +import sys +import tempfile + +os.environ.setdefault("BOT_TOKEN", "123456789:AAHtestTokenPlaceholder_ABCDEFGHIJKLMNOP") +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from aiogram.types import Chat, Message, PhotoSize, User + +logging.disable(logging.CRITICAL) + +CHAT = Chat(id=500, type="private") +USER = User(id=111, is_bot=False, first_name="Owner") +NOW = datetime.datetime.now() + + +def load(path: str, **env): + for key in ("PUBLIC_MODE", "REQUIRED_CHATS", "ALLOWED_USERS", "ADMIN_USERS"): + os.environ.pop(key, None) + os.environ["DB_PATH"] = path + os.environ.update(env) + import db, bot + importlib.reload(db) + B = importlib.reload(bot) + B.db = db + return db, B + + +def photo_msg(message_id: int, caption: str | None = None, group: str | None = None) -> Message: + photo = PhotoSize(file_id=f"F{message_id}", file_unique_id=f"U{message_id}", width=1, height=1) + return Message(message_id=message_id, date=NOW, chat=CHAT, from_user=USER, + photo=[photo], caption=caption, media_group_id=group) + + +# ── database ───────────────────────────────────────────────────────────────── + +async def test_roundtrip(db): + item_id = await db.add_vault_item(111, 500, [12, 10, 11], "A trip to the sea", ["sea"]) + item = await db.get_vault_item(111, item_id) + assert db.unpack_message_ids(item["message_ids"]) == [12, 10, 11] + assert item["caption"] == "A trip to the sea" + print(" a saved post round-trips with its message ids and caption") + + +async def test_search_by_keyword_and_caption(db): + await db.add_vault_item(111, 500, [20], "Recipe for lasagne", ["food"]) + assert len(await db.search_vault(111, "food")) == 1, "keyword search" + assert len(await db.search_vault(111, "lasagne")) == 1, "caption search" + assert len(await db.search_vault(111, "lasagne")) == 1 + assert await db.search_vault(111, "nonsense") == [] + print(" posts are findable by keyword and by words in the original caption") + + +async def test_vault_is_private_to_its_owner(db): + mine = await db.add_vault_item(111, 500, [30], "private note", ["secret"]) + + assert await db.get_vault_item(222, mine) is None, "another user fetched my post" + assert await db.search_vault(222, "secret") == [], "another user searched my vault" + assert await db.delete_vault_item(222, mine) is False, "another user deleted my post" + assert await db.get_vault_item(111, mine) is not None, "owner lost their own post" + print(" one user cannot read, search or delete another user's vault") + + +async def test_delete(db): + item_id = await db.add_vault_item(111, 500, [40], "", ["bye"]) + assert await db.delete_vault_item(111, item_id) is True + assert await db.get_vault_item(111, item_id) is None + assert await db.delete_vault_item(111, item_id) is False + print(" deleting removes the post and is safe to repeat") + + +# ── album grouping ─────────────────────────────────────────────────────────── + +async def test_album_becomes_one_post(B): + B.ALBUM_DEBOUNCE = 0.05 + B._album_buffers.clear() + B._pending_forwards.clear() + sent: list = [] + + async def send_message(chat_id, text, **kw): + sent.append(text) + B.bot.send_message = send_message + + # Telegram delivers an album as separate messages; only one has the caption. + for msg in (photo_msg(10, caption="beach day", group="g1"), + photo_msg(11, group="g1"), + photo_msg(12, group="g1")): + await B.on_forwarded(msg, None) + await asyncio.sleep(0.3) + + assert len(sent) == 1, f"an album should prompt once, got {len(sent)} prompts" + assert len(B._pending_forwards) == 1 + record = next(iter(B._pending_forwards.values())) + assert record["message_ids"] == [10, 11, 12] + assert record["caption"] == "beach day", "the one caption in the album was lost" + assert len(record["media"]) == 3, "every item in the album should be extractable" + print(" an album is buffered into one post, keeping its caption and all media") + + +async def test_single_forward_prompts_immediately(B): + B._album_buffers.clear() + B._pending_forwards.clear() + sent: list = [] + + async def send_message(chat_id, text, **kw): + sent.append(text) + B.bot.send_message = send_message + + await B.on_forwarded(photo_msg(20, caption="just one"), None) + assert len(sent) == 1, "a lone forward should not wait for an album" + record = next(iter(B._pending_forwards.values())) + assert record["message_ids"] == [20] + print(" a single forwarded post prompts straight away") + + +async def test_pending_prompts_are_bounded(B): + B._pending_forwards.clear() + B.bot.send_message = lambda *a, **kw: asyncio.sleep(0) + + for i in range(B.PENDING_FORWARD_LIMIT + 25): + await B.prompt_forward_choice( + {"user_id": 1, "chat_id": 500, "message_ids": [i], "caption": "", "media": []} + ) + assert len(B._pending_forwards) <= B.PENDING_FORWARD_LIMIT + print(" abandoned prompts cannot grow without bound") + + +async def test_retrieval_forwards_ascending_ids(db, B): + item_id = await db.add_vault_item(111, 500, [12, 10, 11], "out of order", ["x"]) + forwarded: dict = {} + + async def forward_messages(chat_id, from_chat_id, message_ids): + forwarded.update(chat_id=chat_id, from_chat_id=from_chat_id, message_ids=message_ids) + B.bot.forward_messages = forward_messages + + class Cb: + data = f"vault:get:{item_id}" + from_user = type("U", (), {"id": 111})() + message = type("M", (), {"chat": CHAT})() + async def answer(self, *a, **kw): pass + + await B.on_vault_action(Cb()) + # Telegram requires strictly increasing ids, and the album must stay whole. + assert forwarded["message_ids"] == [10, 11, 12], forwarded + assert forwarded["from_chat_id"] == 500 + print(" retrieval forwards the whole post with ids in ascending order") + + +async def main(): + with tempfile.TemporaryDirectory() as tmp: + db, B = load(os.path.join(tmp, "vault.db"), ADMIN_USERS="111") + await db.init_db() + + for test in (test_roundtrip, test_search_by_keyword_and_caption, + test_vault_is_private_to_its_owner, test_delete): + await test(db) + for test in (test_album_becomes_one_post, test_single_forward_prompts_immediately, + test_pending_prompts_are_bounded): + await test(B) + await test_retrieval_forwards_ascending_ids(db, B) + print("\nAll vault tests passed.") + + +if __name__ == "__main__": + asyncio.run(main())