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