Commit 9955c39
Changed files (4)
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