main
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3import re
4
5from loguru import logger
6from pyrogram.client import Client
7from pyrogram.types import Message, ReplyParameters
8
9from config import DB, FAVORITE
10from database.r2 import get_cf_r2, set_cf_r2
11from messages.parser import parse_msg
12from messages.sender import send2tg
13from messages.utils import delete_message, equal_prefix, set_reaction, startswith_prefix
14from utils import slim_cid, strings_list, to_int
15
16HELP = f"""⭐️**收藏消息**
17- `{FAVORITE.SAVE_PREFIX} keyword` : 保存消息为"keyword"到收藏夹
18- `{FAVORITE.SEND_PREFIX} keyword` : 从收藏夹发送"keyword"的消息
19
20保存同一条消息为多个关键词时, 需使用逗号分隔
21"""
22
23
24async def save_favorite(client: Client, message: Message, **kwargs):
25 """Save message to Favorites."""
26 # send docs if message == "/save"
27 if equal_prefix(message.text, prefix=FAVORITE.SAVE_PREFIX):
28 await send2tg(client, message, texts=HELP, **kwargs)
29 return
30 if not message.reply_to_message:
31 return
32 info = parse_msg(message, silent=True)
33 if not startswith_prefix(info["text"], prefix=FAVORITE.SAVE_PREFIX):
34 return
35 # check if user can use this command
36 if str(info["uid"]) not in strings_list(FAVORITE.TIDS_ALLOW_SAVE):
37 await message.reply(f"👤{info['full_name']}无权使用【保存收藏】功能\n🆔{info['uid']}", quote=True)
38 return
39 if not all([DB.CF_R2_ENABLED, FAVORITE.BACKUP_CHAT]):
40 await message.reply("【保存收藏】功能缺失必要参数设置", quote=True)
41 return
42
43 await set_reaction(client, message, reaction="👌")
44 keyword = info["text"].removeprefix(FAVORITE.SAVE_PREFIX).strip()
45 is_force_save = keyword.startswith(("!", "！")) # noqa: RUF001
46 keyword = keyword.removeprefix("!").removeprefix("！") # noqa: RUF001
47 msg_to_save = message.reply_to_message
48
49 # custom chat_id and message_id
50 if "cid=" in keyword and "mid=" in keyword:
51 cid = msg_to_save.chat.id
52 mid = msg_to_save.id
53 if matched := re.search(r"cid=(-?\w+)", keyword, re.IGNORECASE):
54 cid = to_int(matched.group(1))
55 keyword = re.sub(r"cid=(-?\w+)", "", keyword).strip()
56 if matched := re.search(r"mid=(\d+)", keyword, re.IGNORECASE):
57 mid = to_int(matched.group(1))
58 keyword = re.sub(r"mid=(\d+)", "", keyword).strip()
59 msg: Message = await client.get_messages(cid, mid) # type: ignore
60 if isinstance(msg, Message):
61 msg_to_save = msg
62 else:
63 await send2tg(client, message, texts=f"❌获取此消息失败, 可能是消息已被删除或者无访问权限.\nChatID={cid}, MessageID={mid}", **kwargs)
64 return
65
66 if not keyword:
67 return
68 keywords = strings_list(keyword.replace("，", ",")) # noqa: RUF001
69 # check if key is existed
70 for key in keywords:
71 if not is_force_save and await get_cf_r2(FAVORITE.R2_PREFIX + key):
72 await send2tg(client, message, texts=f"⚠️【{key}】已存在\n🔄请使用以下命令覆盖收藏:\n`{FAVORITE.SAVE_PREFIX}! {keyword}`", **kwargs)
73 return
74
75 # forward message to backup chat
76 if msg_to_save.media_group_id:
77 messages = await client.get_media_group(msg_to_save.chat.id, msg_to_save.id)
78 save_msgs: list[Message] = await client.forward_messages(
79 chat_id=to_int(FAVORITE.BACKUP_CHAT),
80 from_chat_id=msg_to_save.chat.id,
81 message_ids=[m.id for m in messages],
82 ) # type: ignore
83 save_msg = save_msgs[0]
84 else:
85 save_msg: Message = await msg_to_save.forward(to_int(FAVORITE.BACKUP_CHAT)) # type: ignore
86 metainfo = " ".join([f"#{key}" for key in keywords]) + f"\nhttps://t.me/c/{slim_cid(msg_to_save.chat.id)}/{msg_to_save.id}"
87 await save_msg.reply(metainfo, quote=True)
88 save_info = parse_msg(save_msg, silent=True, use_cache=False)
89 # only allow standard type
90 keys_to_remove = []
91 for k, v in save_info.items():
92 if not isinstance(v, (str, int, float, list, dict)):
93 keys_to_remove.append(k)
94 [save_info.pop(k) for k in keys_to_remove]
95 for key in keywords:
96 if not await set_cf_r2(FAVORITE.R2_PREFIX + key, save_info, silent=True):
97 logger.error(f"保存【{key}】到收藏失败")
98 await set_reaction(client, message, reaction="💔")
99 return
100 logger.success(f"【{key}】已保存到收藏")
101 await set_reaction(client, message, reaction="✍️")
102 await delete_message(message)
103
104
105async def send_favorite(client: Client, message: Message, **kwargs):
106 """Send message from Favorites."""
107 # send docs if message == "/fav"
108 if equal_prefix(message.text, prefix=FAVORITE.SEND_PREFIX):
109 await send2tg(client, message, texts=HELP, **kwargs)
110 return
111 info = parse_msg(message, silent=True)
112 if not startswith_prefix(info["text"], prefix=FAVORITE.SEND_PREFIX):
113 return
114 # check if user can use this command
115 if FAVORITE.TIDS_ALLOW_SEND != "all" and str(info["uid"]) not in strings_list(FAVORITE.TIDS_ALLOW_SEND):
116 await message.reply(f"👤{info['full_name']}无权使用【发送收藏】功能\n🆔{info['uid']}", quote=True)
117 return
118 if not all([DB.CF_R2_ENABLED]):
119 await message.reply("【发送收藏】功能缺失必要参数设置", quote=True)
120 return
121 keyword: str = info["text"].removeprefix(FAVORITE.SEND_PREFIX).strip()
122 if not keyword:
123 return
124
125 await set_reaction(client, message, reaction="👌")
126 save_key = f"{FAVORITE.R2_PREFIX}{keyword}"
127 # check if key is existed
128 save_info = await get_cf_r2(save_key)
129 if not save_info:
130 await send2tg(client, message, texts=f"【{keyword}】不存在", **kwargs)
131 await set_reaction(client, message, reaction="💔")
132 return
133 reply_message_id = message.reply_to_message.id if message.reply_to_message else None
134 reply = ReplyParameters(message_id=reply_message_id) # type: ignore
135 if save_info.get("media_group_id"):
136 succ = await client.copy_media_group(chat_id=info["cid"], from_chat_id=save_info["cid"], message_id=save_info["mid"], reply_parameters=reply)
137 else:
138 succ = await client.copy_message(chat_id=info["cid"], from_chat_id=save_info["cid"], message_id=save_info["mid"], reply_parameters=reply)
139 if isinstance(succ, Message) or (isinstance(succ, list) and all(isinstance(s, Message) for s in succ)):
140 await set_reaction(client, message)
141 else:
142 await set_reaction(client, message, reaction="💔")
143 await delete_message(message)