🚑️🎨 修复词库传入字符错误, 优化代码

2 files changed, 36 insertions, 93 deletions

diff --git a/ATRI/plugins/thesaurus/__init__.py b/ATRI/plugins/thesaurus/__init__.py
index 813e9d7..0409008 100644
--- a/ATRI/plugins/thesaurus/__init__.py
+++ b/ATRI/plugins/thesaurus/__init__.py
@@ -4,6 +4,7 @@ from datetime import timedelta, timezone as tz
 from nonebot.matcher import Matcher
 from nonebot.params import ArgPlainText, CommandArg, ArgStr
 from nonebot.adapters.onebot.v11 import Bot, Message, MessageEvent, GroupMessageEvent
+from nonebot.adapters.onebot.v11.utils import unescape
 
 from ATRI.service import Service
 from ATRI.message import MessageBuilder
@@ -64,8 +65,8 @@ async def _add_normal_item(
     if item_t not in type_list:
         await add_item.finish("该类型不支持 (全匹配、模糊匹配、正则)\n请重新提交.")
 
-    q_checker = MessageChecker(item_q).check_cq_code
-    a_checker = MessageChecker(item_a).check_cq_code
+    q_checker = MessageChecker(unescape(item_q)).check_cq_code
+    a_checker = MessageChecker(unescape(item_a)).check_cq_code
     if not q_checker or not a_checker:
         await add_item.finish("请不要尝试注入！")
 
@@ -79,12 +80,12 @@ async def _add_normal_item(
     )
     operator = operator_info.get("card", "unknown")
     item_id = gen_random_str(6)
-    ans = item_a.split(",,")
+    ans = unescape(item_a).split(",,")
 
     result = await tm.add_item(
         item_id,
         False,
-        item_q,
+        unescape(item_q),
         ans,
         need_at,
         item_t,
@@ -142,8 +143,8 @@ async def _add_group_item(
     if item_t not in type_list:
         await add_item_as_group_admin.finish("该类型不支持 (全匹配、模糊匹配、正则)\n请重新提交.")
 
-    q_checker = MessageChecker(item_q).check_cq_code
-    a_checker = MessageChecker(item_a).check_cq_code
+    q_checker = MessageChecker(unescape(item_q)).check_cq_code
+    a_checker = MessageChecker(unescape(item_a)).check_cq_code
     if not q_checker or not a_checker:
         await add_item_as_group_admin.finish("请不要尝试注入！")
 
@@ -157,12 +158,12 @@ async def _add_group_item(
     )
     operator = operator_info.get("card", "unknown")
     item_id = gen_random_str(6)
-    ans = item_a.split(",,")
+    ans = unescape(item_a).split(",,")
 
     result = await tm.add_item(
         item_id,
         True,
-        item_q,
+        unescape(item_q),
         ans,
         need_at,
         item_t,
@@ -224,10 +225,10 @@ async def _add_global_item(
     operator = "MASTER"
     opeartor_id = event.user_id
     item_id = gen_random_str(6)
-    ans = item_a.split(",,")
+    ans = unescape(item_a).split(",,")
 
     result = await tm.add_item(
-        item_id, True, item_q, ans, need_at, item_t, 0, operator, opeartor_id, 0, list()
+        item_id, True, unescape(item_q), ans, need_at, item_t, 0, operator, opeartor_id, 0, list()
     )
     await add_item_for_global.finish(result)
 
diff --git a/ATRI/plugins/thesaurus/listener.py b/ATRI/plugins/thesaurus/listener.py
index 4e37f90..5c801e8 100644
--- a/ATRI/plugins/thesaurus/listener.py
+++ b/ATRI/plugins/thesaurus/listener.py
@@ -74,100 +74,42 @@ def init_listener():
     )
 
 
-main_listener = thes_listener.on_message(
-    "词库监听器", "监听所有消息判断是否满足触发词条条件", priority=4, block=False
-)
-
-
-@main_listener.handle([Cooldown(3)])
-async def _tl_listener(event: MessageEvent):
+async def process_message(event: MessageEvent, query_group_id: int = int()):
     tl = ThesaurusListener()
     msg = event.get_message().extract_plain_text()
 
-    group_id = int()
-    if isinstance(event, GroupMessageEvent):
-        group_id = event.group_id
-
-    query_result = await tl.get_item_list(group_id)
+    query_result = await tl.get_item_list(query_group_id)
     if not query_result:
-        query_result = await tl.get_item_list(int())
-        if not query_result:
-            return
+        return
 
     shuffle(query_result)
 
-    for item in query_result:
-        item_info: ThesaurusStoragor = item
-
+    for i in query_result:
+        item_info: ThesaurusStoragor = i
         if not item_info.result:
             continue
 
-        if item_info.m_type == 1:
-            if item_info.matcher in msg:
-                if item_info.need_at:
-                    if event.is_tome():
-                        await main_listener.finish(Message(choice(item_info.result)))
-                    else:
-                        return
-                else:
-                    await main_listener.finish(Message(choice(item_info.result)))
-        elif item_info.m_type == 2:
-            patt = item_info.matcher
-            if re.findall(patt, msg):
-                if item_info.need_at:
-                    if event.is_tome():
-                        await main_listener.finish(Message(choice(item_info.result)))
-                    else:
-                        return
-                else:
-                    await main_listener.finish(Message(choice(item_info.result)))
-        else:
-            if item_info.matcher == msg:
-                if item_info.need_at:
-                    if event.is_tome():
-                        await main_listener.finish(Message(choice(item_info.result)))
-                    else:
-                        return
-                else:
-                    await main_listener.finish(Message(choice(item_info.result)))
-
-    query_result = await tl.get_item_list(int())
-    if not query_result:
-        return
+        match_type = item_info.m_type
+        matcher = item_info.matcher
+        if item_info.need_at and not event.is_tome():
+            return
+        if match_type == 1 and matcher in msg:
+            await main_listener.finish(Message(choice(item_info.result)))
+        elif match_type == 2 and re.findall(matcher, msg):
+            await main_listener.finish(Message(choice(item_info.result)))
+        elif matcher == msg:
+            await main_listener.finish(Message(choice(item_info.result)))
 
-    shuffle(query_result)
 
-    for item in query_result:
-        item_info: ThesaurusStoragor = item
+main_listener = thes_listener.on_message(
+    "词库监听器", "监听所有消息判断是否满足触发词条条件", priority=4, block=False
+)
+
 
-        if not item_info.result:
-            continue
+@main_listener.handle([Cooldown(3)])
+async def _(event: MessageEvent):
+    if isinstance(event, GroupMessageEvent):
+        group_id = event.group_id
+        await process_message(event, group_id)
 
-        if item_info.m_type == 1:
-            if item_info.matcher in msg:
-                if item_info.need_at:
-                    if event.is_tome():
-                        await main_listener.finish(Message(choice(item_info.result)))
-                    else:
-                        return
-                else:
-                    await main_listener.finish(Message(choice(item_info.result)))
-        elif item_info.m_type == 2:
-            patt = item_info.matcher
-            if re.findall(patt, msg):
-                if item_info.need_at:
-                    if event.is_tome():
-                        await main_listener.finish(Message(choice(item_info.result)))
-                    else:
-                        return
-                else:
-                    await main_listener.finish(Message(choice(item_info.result)))
-        else:
-            if item_info.matcher == msg:
-                if item_info.need_at:
-                    if event.is_tome():
-                        await main_listener.finish(Message(choice(item_info.result)))
-                    else:
-                        return
-                else:
-                    await main_listener.finish(Message(choice(item_info.result)))
+    await process_message(event)