From 21b956b98398a78cf8568a1e47b98753758ad068 Mon Sep 17 00:00:00 2001
From: zhayujie <yjzha1996@163.com>
Date: Mon, 16 Oct 2023 16:44:06 +0800
Subject: [PATCH] fix: mj open auth bug

---
 plugins/linkai/linkai.py     | 26 ++++----------------------
 plugins/linkai/midjourney.py |  4 ++++
 plugins/linkai/utils.py      | 28 ++++++++++++++++++++++++++++
 3 files changed, 36 insertions(+), 22 deletions(-)
 create mode 100644 plugins/linkai/utils.py

diff --git a/plugins/linkai/linkai.py b/plugins/linkai/linkai.py
index 39feaf124..e40140ab9 100644
--- a/plugins/linkai/linkai.py
+++ b/plugins/linkai/linkai.py
@@ -1,7 +1,6 @@
 import plugins
 from bridge.context import ContextType
 from bridge.reply import Reply, ReplyType
-from config import global_config
 from plugins import *
 from .midjourney import MJBot
 from .summary import LinkSummary
@@ -9,7 +8,7 @@
 from common.expired_dict import ExpiredDict
 from common import const
 import os
-
+from .utils import Util
 
 @plugins.register(
     name="linkai",
@@ -129,7 +128,7 @@ def _process_admin_cmd(self, e_context: EventContext):
 
         if len(cmd) == 2 and (cmd[1] == "open" or cmd[1] == "close"):
             # 知识库开关指令
-            if not _is_admin(e_context):
+            if not Util.is_admin(e_context):
                 _set_reply_text("需要管理员权限执行", e_context, level=ReplyType.ERROR)
                 return
             is_open = True
@@ -147,7 +146,7 @@ def _process_admin_cmd(self, e_context: EventContext):
             if not context.kwargs.get("isgroup"):
                 _set_reply_text("该指令需在群聊中使用", e_context, level=ReplyType.ERROR)
                 return
-            if not _is_admin(e_context):
+            if not Util.is_admin(e_context):
                 _set_reply_text("需要管理员权限执行", e_context, level=ReplyType.ERROR)
                 return
             app_code = cmd[2]
@@ -164,7 +163,7 @@ def _process_admin_cmd(self, e_context: EventContext):
 
         if len(cmd) == 3 and cmd[1] == "sum" and (cmd[2] == "open" or cmd[2] == "close"):
             # 知识库开关指令
-            if not _is_admin(e_context):
+            if not Util.is_admin(e_context):
                 _set_reply_text("需要管理员权限执行", e_context, level=ReplyType.ERROR)
                 return
             is_open = True
@@ -253,23 +252,6 @@ def _send_info(e_context: EventContext, content: str):
     channel = e_context["channel"]
     channel.send(reply, e_context["context"])
 
-# 静态方法
-def _is_admin(e_context: EventContext) -> bool:
-    """
-    判断消息是否由管理员用户发送
-    :param e_context: 消息上下文
-    :return: True: 是, False: 否
-    """
-    context = e_context["context"]
-    if context["isgroup"]:
-        actual_user_id= context.kwargs.get("msg").actual_user_id
-        for admin_user in global_config["admin_users"]:
-            if actual_user_id and actual_user_id in admin_user:
-                return True
-        return False
-    else:
-        return context["receiver"] in global_config["admin_users"]
-
 
 def _find_user_id(context):
     if context["isgroup"]:
diff --git a/plugins/linkai/midjourney.py b/plugins/linkai/midjourney.py
index a944ae5de..76395bd5a 100644
--- a/plugins/linkai/midjourney.py
+++ b/plugins/linkai/midjourney.py
@@ -8,6 +8,7 @@
 import asyncio
 from bridge.context import ContextType
 from plugins import EventContext, EventAction
+from .utils import Util
 
 INVALID_REQUEST = 410
 NOT_FOUND_ORIGIN_IMAGE = 461
@@ -113,6 +114,9 @@ def process_mj_task(self, mj_type: TaskType, e_context: EventContext):
             return
 
         if len(cmd) == 2 and (cmd[1] == "open" or cmd[1] == "close"):
+            if not Util.is_admin(e_context):
+                Util.set_reply_text("需要管理员权限执行", e_context, level=ReplyType.ERROR)
+                return
             # midjourney 开关指令
             is_open = True
             tips_text = "开启"
diff --git a/plugins/linkai/utils.py b/plugins/linkai/utils.py
new file mode 100644
index 000000000..c874cdfb1
--- /dev/null
+++ b/plugins/linkai/utils.py
@@ -0,0 +1,28 @@
+from config import global_config
+from bridge.reply import Reply, ReplyType
+from plugins.event import EventContext, EventAction
+
+
+class Util:
+    @staticmethod
+    def is_admin(e_context: EventContext) -> bool:
+        """
+        判断消息是否由管理员用户发送
+        :param e_context: 消息上下文
+        :return: True: 是, False: 否
+        """
+        context = e_context["context"]
+        if context["isgroup"]:
+            actual_user_id = context.kwargs.get("msg").actual_user_id
+            for admin_user in global_config["admin_users"]:
+                if actual_user_id and actual_user_id in admin_user:
+                    return True
+            return False
+        else:
+            return context["receiver"] in global_config["admin_users"]
+
+    @staticmethod
+    def set_reply_text(content: str, e_context: EventContext, level: ReplyType = ReplyType.ERROR):
+        reply = Reply(level, content)
+        e_context["reply"] = reply
+        e_context.action = EventAction.BREAK_PASS