main
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3import os
4import re
5
6from glom import glom
7from loguru import logger
8from pyrogram.types import Message
9from pyrogram.types.messages_and_media.message import Str
10
11from messages.parser import get_thread_id
12from utils import slim_cid
13
14
15def message_modify(message: Message) -> Message:
16 """This function modifies the message via environment variable before any subsequent processing.
17
18 Args:
19 message (Message): The message to be modified.
20
21 Returns:
22 Message: The modified message.
23 """
24 message = add_prefix_suffix(message)
25 return replace_prefix(message)
26
27
28def add_prefix_suffix(message: Message) -> Message:
29 """This function add prefix and suffix to the message via environment variable.
30
31 Environment Variables:
32 MOD_MSG_ADD_PREFIX_C{CID} (str): Add prefix to message of {CID}
33 MOD_MSG_ADD_SUFFIX_C{CID} (str): Add suffix to message of {CID}
34 MOD_MSG_ADD_PREFIX_C{CID}_T{TID} (str): Add prefix to message of {CID}-{TID}
35 MOD_MSG_ADD_SUFFIX_C{CID}_T{TID} (str): Add suffix to message of {CID}-{TID}
36 MOD_MSG_ADD_PREFIX_U{UID} (str): Add prefix to message of {UID}
37 MOD_MSG_ADD_SUFFIX_U{UID} (str): Add suffix to message of {UID}
38 MOD_MSG_ADD_PREFIX_C{CID}_U{UID} (str): Add prefix to message of {CID} and {UID}
39 MOD_MSG_ADD_SUFFIX_C{CID}_U{UID} (str): Add suffix to message of {CID} and {UID}
40 MOD_MSG_ADD_PREFIX_C{CID}_T{TID}_U{UID} (str): Add prefix to message of {CID}-{TID} and {UID}
41 MOD_MSG_ADD_SUFFIX_C{CID}_T{TID}_U{UID} (str): Add suffix to message of {CID}-{TID} and {UID}
42
43 If there are multiple envvars for the same message, they will be applied in the order of
44 Chat Level -> Chat_Thread Level -> User Level -> Chat_User Level -> Chat_Thread_User Level
45
46 Example:
47 Suppose the message is Message(chat_id=111,from_user=User(id=222), content="hello", message_thread_id=333)
48
49 # Single Environment Variable:
50 MOD_MSG_ADD_PREFIX_C111 = "foo" -> "foo hello" for chat_id=111
51 MOD_MSG_ADD_SUFFIX_C111 = "bar" -> "hello bar" for chat_id=111
52 MOD_MSG_ADD_PREFIX_C111_T333 = "foo" -> "foo hello" for chat_id=111, message_thread_id=333
53 MOD_MSG_ADD_SUFFIX_C111_T333 = "bar" -> "hello bar" for chat_id=111, message_thread_id=333
54 MOD_MSG_ADD_PREFIX_U222 = "foo" -> "foo hello" for user_id=222
55 MOD_MSG_ADD_SUFFIX_U222 = "bar" -> "hello bar" for user_id=222
56 MOD_MSG_ADD_PREFIX_C111_U222 = "foo" -> "foo hello" for chat_id=111 and user_id=222
57 MOD_MSG_ADD_SUFFIX_C111_U222 = "bar" -> "hello bar" for chat_id=111 and user_id=222
58
59 # Multiple Environment Variables:
60 MOD_MSG_ADD_PREFIX_C111 = "foo" + MOD_MSG_ADD_PREFIX_C111_U222 = "bar"
61 -> "bar foo hello" (Chat level prefix `foo`, then Chat_User level prefix `bar`)
62
63 MOD_MSG_ADD_SUFFIX_U111 = "foo" + MOD_MSG_ADD_PREFIX_C111 = "bar"
64 -> "bar hello foo" (Chat level prefix `bar`, then User level suffix `foo`)
65
66
67 Args:
68 message (Message): The message to be modified.
69
70 Returns:
71 Message: The modified message.
72 """
73 uid = glom(message, "from_user.id", default=0) or 0
74 cid = glom(message, "chat.id", default=0) or 0
75 tid = get_thread_id(message)
76 cid = slim_cid(cid)
77 texts = str(message.content).strip()
78 # Chat level
79 suffix = os.getenv(f"MOD_MSG_ADD_SUFFIX_C{cid}")
80 if prefix := os.getenv(f"MOD_MSG_ADD_PREFIX_C{cid}"):
81 texts = f"{prefix} {texts}"
82 if suffix := os.getenv(f"MOD_MSG_ADD_SUFFIX_C{cid}"):
83 texts = f"{texts} {suffix}"
84
85 # Chat_Thread level
86 suffix = os.getenv(f"MOD_MSG_ADD_SUFFIX_C{cid}_T{tid}")
87 if prefix := os.getenv(f"MOD_MSG_ADD_PREFIX_C{cid}_T{tid}"):
88 texts = f"{prefix} {texts}"
89 if suffix := os.getenv(f"MOD_MSG_ADD_SUFFIX_C{cid}_T{tid}"):
90 texts = f"{texts} {suffix}"
91
92 # User level
93 if prefix := os.getenv(f"MOD_MSG_ADD_PREFIX_U{uid}"):
94 texts = f"{prefix} {texts}"
95 if suffix := os.getenv(f"MOD_MSG_ADD_SUFFIX_U{uid}"):
96 texts = f"{texts} {suffix}"
97
98 # Chat_User level
99 if prefix := os.getenv(f"MOD_MSG_ADD_PREFIX_C{cid}_U{uid}"):
100 texts = f"{prefix} {texts}"
101 if suffix := os.getenv(f"MOD_MSG_ADD_SUFFIX_C{cid}_U{uid}"):
102 texts = f"{texts} {suffix}"
103
104 # Chat_Thread_User level
105 if prefix := os.getenv(f"MOD_MSG_ADD_PREFIX_C{cid}_T{tid}_U{uid}"):
106 texts = f"{prefix} {texts}"
107 if suffix := os.getenv(f"MOD_MSG_ADD_SUFFIX_C{cid}_T{tid}_U{uid}"):
108 texts = f"{texts} {suffix}"
109
110 if texts.strip() != str(message.content).strip(): # If the message is modified
111 if message.text:
112 message.text = Str(texts)
113 else:
114 message.caption = Str(texts)
115 return message
116
117
118def replace_prefix(message: Message) -> Message:
119 """This function replace prefix of the message via environment variables.
120
121 Environment Variables:
122 MOD_MSG_REPLACE_PREFIX_GLOBAL_{prefix} (str): Replace prefix to all messages (prefix is case-sensitive)
123 MOD_MSG_REPLACE_PREFIX_C{CID}_{prefix} (str): Replace prefix to message of {CID} (prefix is case-sensitive)
124 MOD_MSG_REPLACE_PREFIX_CT{CID}_{TID}_{prefix} (str): Replace prefix to message of {CID}-{TID} (prefix is case-sensitive)
125 MOD_MSG_REPLACE_PREFIX_U{UID}_{prefix} (str): Replace prefix to message of {UID} (prefix is case-sensitive)
126 MOD_MSG_REPLACE_PREFIX_CU{CID}_{UID}_{prefix} (str): Replace prefix to message of {CID} and {UID} (prefix is case-sensitive)
127 MOD_MSG_REPLACE_PREFIX_CTU{CID}_{TID}_{UID}_{prefix} (str): Replace prefix to message of {CID}-{TID} and {UID} (prefix is case-sensitive)
128
129 If there are multiple envvars for the same message, they will be applied in the order of
130 Global Level -> Chat Level -> User Level -> Chat_User Level
131
132 Example:
133 Suppose the message is Message(chat_id=111, from_user=User(id=222), content="this is a sample msg")
134 # Single Environment Variable:
135 MOD_MSG_REPLACE_PREFIX_GLOBAL_this = "foo" -> "foo is a sample msg" for all messages
136 MOD_MSG_REPLACE_PREFIX_C111_this = "foo" -> "foo is a sample msg" for chat_id=111
137 MOD_MSG_REPLACE_PREFIX_U222_this__SPACE__ = "foo" -> "foois a sample msg" for user_id=222
138 MOD_MSG_REPLACE_PREFIX_CU111_222_this__SPACE__is = "foo" -> "fooa sample msg" for chat_id=111 and user_id=222
139
140 # Multiple Environment Variables:
141 MOD_MSG_REPLACE_PREFIX_GLOBAL_this = "foo" + MOD_MSG_REPLACE_PREFIX_C111_foo = "bar"
142 -> "bar is a sample msg" (Global level `this` -> `foo`, then Chat level `foo` -> `bar`)
143
144 Args:
145 message (Message): The message to be modified.
146
147 Returns:
148 Message: The modified message.
149 """
150 uid = glom(message, "from_user.id", default=0) or 0
151 cid = glom(message, "chat.id", default=0) or 0
152 tid = get_thread_id(message)
153 cid = slim_cid(cid)
154 texts = str(message.content).strip()
155
156 # Global level
157 global_envs = [x for x in os.environ if x.upper().startswith("MOD_MSG_REPLACE_PREFIX_GLOBAL_")]
158 for env in sorted(global_envs, key=lambda x: len(x), reverse=True): # sorted by length in descending order
159 env_var = os.environ[env]
160 prefix = env[len("MOD_MSG_REPLACE_PREFIX_GLOBAL_") :]
161 prefix = escape_strings(prefix)
162 env_var = escape_strings(env_var)
163 if texts.startswith(prefix):
164 texts = texts.replace(prefix, env_var, count=1)
165 logger.warning(f"Global level prefix `{prefix} = {env_var}` -> {texts}")
166
167 # Chat level
168 chat_envs = [x for x in os.environ if x.upper().startswith(f"MOD_MSG_REPLACE_PREFIX_C{cid}_")]
169 for env in sorted(chat_envs, key=lambda x: len(x), reverse=True):
170 env_var = os.environ[env]
171 prefix = env[len(f"MOD_MSG_REPLACE_PREFIX_C{cid}_") :]
172 prefix = escape_strings(prefix)
173 env_var = escape_strings(env_var)
174 if texts.startswith(prefix):
175 texts = texts.replace(prefix, env_var, count=1)
176 logger.warning(f"Chat ({cid}) level prefix `{prefix} = {env_var}` -> {texts}")
177
178 # Chat_Thread level
179 chat_thread_envs = [x for x in os.environ if x.upper().startswith(f"MOD_MSG_REPLACE_PREFIX_CT{cid}_{tid}_")]
180 for env in sorted(chat_thread_envs, key=lambda x: len(x), reverse=True):
181 env_var = os.environ[env]
182 prefix = env[len(f"MOD_MSG_REPLACE_PREFIX_CT{cid}_{tid}_") :]
183 prefix = escape_strings(prefix)
184 env_var = escape_strings(env_var)
185 if texts.startswith(prefix):
186 texts = texts.replace(prefix, env_var, count=1)
187 logger.warning(f"Chat ({cid}), Thread ({tid}) level prefix `{prefix} = {env_var}` -> {texts}")
188
189 # User level
190 user_envs = [x for x in os.environ if x.upper().startswith(f"MOD_MSG_REPLACE_PREFIX_U{uid}_")]
191 for env in sorted(user_envs, key=lambda x: len(x), reverse=True):
192 env_var = os.environ[env]
193 prefix = env[len(f"MOD_MSG_REPLACE_PREFIX_U{uid}_") :]
194 prefix = escape_strings(prefix)
195 env_var = escape_strings(env_var)
196 if texts.startswith(prefix):
197 texts = texts.replace(prefix, env_var, count=1)
198 logger.warning(f"User ({uid}) level prefix `{prefix} = {env_var}` -> {texts}")
199
200 # Chat_User level
201 chat_user_envs = [x for x in os.environ if x.upper().startswith(f"MOD_MSG_REPLACE_PREFIX_CU{cid}_{uid}_")]
202 for env in sorted(chat_user_envs, key=lambda x: len(x), reverse=True):
203 env_var = os.environ[env]
204 prefix = env[len(f"MOD_MSG_REPLACE_PREFIX_CU{cid}_{uid}_") :]
205 prefix = escape_strings(prefix)
206 env_var = escape_strings(env_var)
207 if texts.startswith(prefix):
208 texts = texts.replace(prefix, env_var, count=1)
209 logger.warning(f"Chat ({cid}), User ({uid}) level prefix `{prefix} = {env_var}` -> {texts}")
210
211 # Chat_Thread_User level
212 chat_thread_user_envs = [x for x in os.environ if x.upper().startswith(f"MOD_MSG_REPLACE_PREFIX_CTU{cid}_{tid}_{uid}_")]
213 for env in sorted(chat_thread_user_envs, key=lambda x: len(x), reverse=True):
214 env_var = os.environ[env]
215 prefix = env[len(f"MOD_MSG_REPLACE_PREFIX_CTU{cid}_{tid}_{uid}_") :]
216 prefix = escape_strings(prefix)
217 env_var = escape_strings(env_var)
218 if texts.startswith(prefix):
219 texts = texts.replace(prefix, env_var, count=1)
220 logger.warning(f"Chat ({cid}), Thread ({tid}), User ({uid}) level prefix `{prefix} = {env_var}` -> {texts}")
221
222 if texts != str(message.content).strip(): # If the message is modified
223 if message.text:
224 message.text = Str(texts)
225 else:
226 message.caption = Str(texts)
227 return message
228
229
230STRING_MAP = {
231 "__DASH__": "-",
232 "__SPACE__": " ",
233 "__COLON__": ":",
234 "__PERIOD__": ".",
235 "__COMMA__": ",",
236 "__SEMICOLON__": ";",
237 "__EXCLAMATION__": "!",
238 "__QUESTION__": "?",
239 "__LEFT_PARENTHESIS__": "(",
240 "__RIGHT_PARENTHESIS__": ")",
241 "__LEFT_BRACKET__": "[",
242 "__RIGHT_BRACKET__": "]",
243 "__LEFT_BRACE__": "{",
244 "__RIGHT_BRACE__": "}",
245 "__SLASH__": "/",
246 "__ASTERISK__": "*",
247 "__POUND__": "#",
248 "__NEWLINE__": "\n",
249 "__AT__": "@",
250 "__DOLLAR__": "$",
251 "__PERCENT__": "%",
252 "__CARET__": "^",
253 "__AMPERSAND__": "&",
254 "__UNDERSCORE__": "_",
255 "__PLUS__": "+",
256 "__EQUALS__": "=",
257 "__PIPE__": "|",
258 "__BACKSLASH__": "\\",
259 "__BACKTICK__": "`",
260 "__QUOTE__": "'",
261 "__DOUBLE_QUOTE__": '"',
262 "__LESS_THAN__": "<",
263 "__GREATER_THAN__": ">",
264 "__TILDE__": "~",
265}
266STRING_RE_PATTERN = re.compile("|".join(re.escape(k) for k in STRING_MAP))
267
268
269def escape_strings(s: str) -> str:
270 """Escape some special characters.
271
272 "foo__DASH__bar" -> "foo-bar"
273 "http__COLON____SLASH____SLASH__www__DOT__example__DOT__com" -> "http://www.example.com"
274 """
275 if not isinstance(s, str) or "__" not in s:
276 return s
277 return STRING_RE_PATTERN.sub(lambda m: STRING_MAP[m.group(0)], s)
278
279
280def parse_kwargs(message: Message) -> tuple[Message, dict]:
281 """Parse the kwargs from the message text.
282
283 Currently, three formats are supported:
284 1. #no_xx: kwargs["xx"] = False
285 2. #with_xx: kwargs["xx"] = True
286 3. #set_xx=var: kwargs["xx"] = var
287
288 Example:
289 Suppose the message is Message(text="sample texts #no_ytdlp_send_video #set_douyin_provider=tikhub")
290 Then the kwargs will be:
291 kwargs = {
292 "ytdlp_send_video": False,
293 "douyin_provider": "tikhub",
294 }
295 """
296 if not isinstance(message, Message):
297 return message, {}
298
299 if not message.content:
300 return message, {}
301
302 kwargs = {}
303 texts = str(message.content)
304 # Pattern 1: #no_xx -> kwargs["xx"] = False
305 for match in re.findall(r"#no_(\w+)", texts, flags=re.IGNORECASE):
306 kwargs[match.lower()] = False
307 texts = re.sub(rf"#no_{match}", "", texts, flags=re.IGNORECASE)
308
309 # Pattern 2: #with_xx -> kwargs["xx"] = True
310 for match in re.findall(r"#with_(\w+)", texts, re.IGNORECASE):
311 kwargs[match.lower()] = True
312 texts = re.sub(rf"#with_{match}", "", texts, flags=re.IGNORECASE)
313
314 # Pattern 3: #set_xx=var -> kwargs["xx"] = var
315 for match in re.findall(r"#set_(\w+)=([^#\n]+)", texts, re.IGNORECASE):
316 key, value = match
317 texts = re.sub(rf"#set_{key}={value}", "", texts, flags=re.IGNORECASE)
318 if str(value).strip().lower() in ["none", "null"]:
319 value = None
320 kwargs[key.lower()] = value.strip() if value is not None else None
321 if kwargs:
322 logger.info(f"🔧手动设置参数: {kwargs}")
323 if message.text:
324 message.text = Str(texts)
325 else:
326 message.caption = Str(texts)
327 return message, kwargs