Commit 9955c39

benny-dou <60535774+benny-dou@users.noreply.github.com>
2025-07-03 10:07:38
feat(favorite): add `/fav` and `/save` commands
1 parent b0b2ca7
src/others/favorite.py
@@ -0,0 +1,130 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+import contextlib
+
+from loguru import logger
+from pyrogram.client import Client
+from pyrogram.types import Message
+
+from config import DB, FAVORITE
+from database.r2 import get_cf_r2, set_cf_r2
+from messages.parser import parse_msg
+from messages.sender import send2tg
+from messages.utils import equal_prefix, startswith_prefix
+from utils import strings_list, to_int
+
+HELP = f"""⭐️**收藏消息**
+- `{FAVORITE.SAVE_PREFIX} keyword` : 保存消息为"keyword"到收藏夹
+- `{FAVORITE.SEND_PREFIX} keyword` : 从收藏夹发送"keyword"的消息
+- `{FAVORITE.SEND_PREFIX} fwd keyword` : 转发原始消息
+
+保存同一条消息为多个关键词时, 需使用逗号分隔
+"""
+PRESERVED_KEYS = {"fwd"}  # these keys are preserved
+
+
+async def save_favorite(client: Client, message: Message, **kwargs):
+    """Save message to Favorites."""
+    # send docs if message == "/save"
+    if equal_prefix(message.text, prefix=FAVORITE.SAVE_PREFIX):
+        await send2tg(client, message, texts=HELP, **kwargs)
+        return
+    if not message.reply_to_message:
+        return
+    info = parse_msg(message, silent=True)
+    if not startswith_prefix(info["text"], prefix=FAVORITE.SAVE_PREFIX):
+        return
+    # check if user can use this command
+    if str(info["uid"]) not in strings_list(FAVORITE.TIDS_ALLOW_SAVE):
+        await message.reply(f"👤{info['full_name']}无权使用【保存收藏】功能\n🆔{info['uid']}", quote=True)
+        return
+    if not all([DB.CF_R2_ENABLED, FAVORITE.BACKUP_CHAT]):
+        await message.reply("【保存收藏】功能缺失必要参数设置", quote=True)
+        return
+
+    keyword = info["text"].removeprefix(FAVORITE.SAVE_PREFIX).strip()
+    is_force_save = keyword.startswith(("!", "!"))  # noqa: RUF001
+    keyword = keyword.removeprefix("!").removeprefix("!")  # noqa: RUF001
+    if not keyword:
+        return
+    if keyword.lower() in PRESERVED_KEYS:
+        await message.reply(f"【{keyword}】为保留关键字, 请使用其他关键词", quote=True)
+        return
+    keywords = strings_list(keyword.replace(",", ","))  # noqa: RUF001
+    # check if key is existed
+    for key in keywords:
+        if not is_force_save and await get_cf_r2(FAVORITE.R2_PREFIX + key):
+            await send2tg(client, message, texts=f"⚠️【{key}】已存在\n🔄请使用以下命令覆盖收藏:\n`{FAVORITE.SAVE_PREFIX}! {keyword}`", **kwargs)
+            return
+
+    # forward message to backup chat
+    reply_msg = message.reply_to_message
+    if reply_msg.media_group_id:
+        messages = await client.get_media_group(reply_msg.chat.id, reply_msg.id)
+        save_msgs: list[Message] = await client.forward_messages(
+            chat_id=to_int(FAVORITE.BACKUP_CHAT),
+            from_chat_id=reply_msg.chat.id,
+            message_ids=[m.id for m in messages],
+        )  # type: ignore
+        save_msg = save_msgs[0]
+    else:
+        save_msg: Message = await reply_msg.forward(to_int(FAVORITE.BACKUP_CHAT))  # type: ignore
+    await save_msg.reply(" ".join([f"#{key}" for key in keywords]), quote=True)
+    save_info = parse_msg(save_msg, silent=True, use_cache=False)
+    # only allow standard type
+    keys_to_remove = []
+    for k, v in save_info.items():
+        if not isinstance(v, (str, int, float, list, dict)):
+            keys_to_remove.append(k)
+    [save_info.pop(k) for k in keys_to_remove]
+    for key in keywords:
+        if not await set_cf_r2(FAVORITE.R2_PREFIX + key, save_info, silent=True):
+            logger.error(f"保存【{key}】到收藏失败")
+            return
+        logger.success(f"【{key}】已保存到收藏")
+    with contextlib.suppress(Exception):
+        await message.delete()
+
+
+async def send_favorite(client: Client, message: Message, **kwargs):
+    """Send message from Favorites."""
+    # send docs if message == "/fav"
+    if equal_prefix(message.text, prefix=FAVORITE.SEND_PREFIX):
+        await send2tg(client, message, texts=HELP, **kwargs)
+        return
+    info = parse_msg(message, silent=True)
+    if not startswith_prefix(info["text"], prefix=FAVORITE.SEND_PREFIX):
+        return
+    # check if user can use this command
+    if FAVORITE.TIDS_ALLOW_SEND != "all" and str(info["uid"]) not in strings_list(FAVORITE.TIDS_ALLOW_SEND):
+        await message.reply(f"👤{info['full_name']}无权使用【发送收藏】功能\n🆔{info['uid']}", quote=True)
+        return
+    if not all([DB.CF_R2_ENABLED]):
+        await message.reply("【发送收藏】功能缺失必要参数设置", quote=True)
+        return
+    keyword: str = info["text"].removeprefix(FAVORITE.SEND_PREFIX).strip()
+    use_forward = keyword.lower().startswith("fwd ")
+    if use_forward:
+        keyword = keyword[4:]
+    if not keyword:
+        return
+
+    save_key = f"{FAVORITE.R2_PREFIX}{keyword}"
+    # check if key is existed
+    save_info = await get_cf_r2(save_key)
+    if not save_info:
+        await send2tg(client, message, texts=f"【{keyword}】不存在", **kwargs)
+        return
+
+    if use_forward:
+        if save_info.get("media_group_id"):
+            messages = await client.get_media_group(save_info["cid"], save_info["mid"])
+            await client.forward_messages(chat_id=info["cid"], from_chat_id=save_info["cid"], message_ids=[m.id for m in messages])
+        else:
+            await client.forward_messages(chat_id=info["cid"], from_chat_id=save_info["cid"], message_ids=save_info["mid"])
+    elif save_info.get("media_group_id"):
+        await client.copy_media_group(chat_id=info["cid"], from_chat_id=save_info["cid"], message_id=save_info["mid"])
+    else:
+        await client.copy_message(chat_id=info["cid"], from_chat_id=save_info["cid"], message_id=save_info["mid"])
+    with contextlib.suppress(Exception):
+        await message.delete()
src/config.py
@@ -63,6 +63,7 @@ class ENABLE:  # see fine-grained permission in `src/permission.py`
     SEND_AS_REPLY = os.getenv("ENABLE_SEND_AS_REPLY", "1").lower() in ["1", "y", "yes", "t", "true", "on"]  # Send as a reply to the original message
     CACHE_PRICE_SYMBOLS = os.getenv("ENABLE_CACHE_PRICE_SYMBOLS", "0").lower() in ["1", "y", "yes", "t", "true", "on"]
     QUERY_DANMU = os.getenv("ENABLE_QUERY_DANMU", "1").lower() in ["1", "y", "yes", "t", "true", "on"]
+    FAVORITE = os.getenv("ENABLE_FAVORITE", "1").lower() in ["1", "y", "yes", "t", "true", "on"]
 
 
 class PREFIX:
@@ -282,6 +283,17 @@ class PODCAST:
     KEEP_LATEST_ENTRIES = int(os.getenv("PODCAST_KEEP_LATEST_ENTRIES", "99999999"))  # keep latest entries
 
 
+class FAVORITE:
+    ENABLE_SEND = os.getenv("ENABLE_FAVORITE_SEND", "1").lower() in ["1", "y", "yes", "t", "true", "on"]
+    ENABLE_SAVE = os.getenv("ENABLE_FAVORITE_SAVE", "1").lower() in ["1", "y", "yes", "t", "true", "on"]
+    SEND_PREFIX = os.getenv("FAVORITE_SEND_PREFIX", "/fav").lower()
+    SAVE_PREFIX = os.getenv("FAVORITE_SAVE_PREFIX", "/save").lower()
+    R2_PREFIX = os.getenv("FAVORITE_R2_PREFIX", "Favorite/")
+    BACKUP_CHAT = os.getenv("FAVORITE_BACKUP_CHAT", "")  # chat id to backup favorite messages
+    TIDS_ALLOW_SEND = os.getenv("FAVORITE_TIDS_ALLOW_SEND", "all").lower()  # or comma separated telegram uids
+    TIDS_ALLOW_SAVE = os.getenv("FAVORITE_TIDS_ALLOW_SAVE", "")  # comma separated telegram uids
+
+
 class GPT:
     """This is for OpenAI compatible API.
 
src/handler.py
@@ -20,6 +20,7 @@ from messages.utils import equal_prefix, startswith_prefix
 from networking import match_social_media_link
 from others.download_external import download_url_in_message
 from others.extract_audio import extract_audio_file
+from others.favorite import save_favorite, send_favorite
 from others.raw_img_file import convert_raw_img_file
 from others.search_google import search_google
 from others.search_ytb import search_youtube
@@ -48,6 +49,7 @@ async def handle_utilities(
     asr: bool = True,
     audio: bool = True,
     danmu: bool = True,
+    favorite: bool = True,
     google: bool = True,
     ocr: bool = True,
     history: bool = True,
@@ -74,6 +76,7 @@ async def handle_utilities(
         asr (bool, optional): Enable ASR. Defaults to True.
         audio (bool, optional): Enable Video -> Audio. Defaults to True.
         danmu (bool, optional): Enable Query Danmu database. Defaults to True.
+        favorite (bool, optional): Enable Send & Save message to favorite. Defaults to True.
         google (bool, optional): Enable Google Search. Defaults to True.
         ytb (bool, optional): Enable YouTube Search. Defaults to True.
         history (bool, optional): Enable History Search. Defaults to True.
@@ -113,6 +116,9 @@ async def handle_utilities(
         await ai_summary(client, message, **kwargs)  # /summary
     if danmu:
         await query_danmu(client, message, **kwargs)  # /danmu
+    if favorite:
+        await save_favorite(client, message, **kwargs)  # /save
+        await send_favorite(client, message, **kwargs)  # /fav
     if raw_img:
         await convert_raw_img_file(client, message, **kwargs)
 
src/permission.py
@@ -119,6 +119,7 @@ def check_service(cid: int | str, ctype: str) -> dict:
         "reddit": True,
         "ytdlp": True,
         "history": True,
+        "favorite": True,
     }
 
     if ctype == "PRIVATE":
@@ -168,6 +169,8 @@ def check_service(cid: int | str, ctype: str) -> dict:
         permission["danmu"] = False
     if not ENABLE.HISTORY:
         permission["history"] = False
+    if not ENABLE.FAVORITE:
+        permission["favorite"] = False
 
     """
     Set specific service