✨ 将cd，block，count限制复原配置文件 (#1662)

zhenxun/builtin_plugins/hooks/_auth_checker.py

@@ -88,7 +88,7 @@ class LimitManage:
     @classmethod
     async def check(
         cls,
-        module: str,
+        module_path: str,
         user_id: str,
         group_id: str | None,
         channel_id: str | None,
@@ -106,17 +106,17 @@ class LimitManage:
         异常:
             IgnoredException: IgnoredException
         """
-        if limit_model := cls.cd_limit.get(module):
+        if limit_model := cls.cd_limit.get(module_path):
             await cls.__check(limit_model, user_id, group_id, channel_id, session)
-        if limit_model := cls.block_limit.get(module):
+        if limit_model := cls.block_limit.get(module_path):
             await cls.__check(limit_model, user_id, group_id, channel_id, session)
-        if limit_model := cls.count_limit.get(module):
+        if limit_model := cls.count_limit.get(module_path):
             await cls.__check(limit_model, user_id, group_id, channel_id, session)
 
     @classmethod
     async def __check(
         cls,
-        limit_model: Limit,
+        limit_model: Limit | None,
         user_id: str,
         group_id: str | None,
         channel_id: str | None,
@@ -291,12 +291,14 @@ class AuthChecker:
         if not group_id:
             group_id = channel_id
             channel_id = None
-        limit_list: list[PluginLimit] = await plugin.plugin_limit.all()  # type: ignore
+        limit_list: list[PluginLimit] = await plugin.plugin_limit.filter(
+            status=True
+        ).all()  # type: ignore
         for limit in limit_list:
             LimitManage.add_limit(limit)
         if user_id:
             await LimitManage.check(
-                plugin.module, user_id, group_id, channel_id, session
+                plugin.module_path, user_id, group_id, channel_id, session
             )
 
     async def auth_plugin(

zhenxun/builtin_plugins/init/init_config.py

@@ -1,16 +1,16 @@
 from pathlib import Path
 
 import nonebot
-from nonebot import get_loaded_plugins
-from nonebot.drivers import Driver
-from nonebot.plugin import Plugin
 from ruamel.yaml import YAML
+from nonebot.plugin import Plugin
+from nonebot.drivers import Driver
+from nonebot import get_loaded_plugins
 from ruamel.yaml.comments import CommentedMap
 
-from zhenxun.configs.config import Config
-from zhenxun.configs.path_config import DATA_PATH
-from zhenxun.configs.utils import RegisterConfig
 from zhenxun.services.log import logger
+from zhenxun.configs.config import Config
+from zhenxun.configs.utils import RegisterConfig
+from zhenxun.configs.path_config import DATA_PATH
 
 _yaml = YAML(pure=True)
 _yaml.allow_unicode = True
@@ -72,15 +72,14 @@ def _generate_simple_config():
                     Config.set_config(module, k, _data[module][k])
                 _tmp_data[module][k] = Config.get_config(module, k)
             except AttributeError as e:
-                raise AttributeError(f"{e}\n" + "可能为config.yaml配置文件填写不规范")
+                raise AttributeError(f"{e}\n可能为config.yaml配置文件填写不规范") from e
     Config.save()
     temp_file = DATA_PATH / "temp_config.yaml"
     # 重新生成简易配置文件
     try:
         with open(temp_file, "w", encoding="utf8") as wf:
-            # yaml.dump(_tmp_data, wf, Dumper=yaml.RoundTripDumper, allow_unicode=True)
             _yaml.dump(_tmp_data, wf)
-        with open(temp_file, "r", encoding="utf8") as rf:
+        with open(temp_file, encoding="utf8") as rf:
             _data = _yaml.load(rf)
         # 添加注释
         for module in _data.keys():
@@ -93,7 +92,7 @@ def _generate_simple_config():
         with SIMPLE_CONFIG_FILE.open("w", encoding="utf8") as wf:
             _yaml.dump(_data, wf)
     except Exception as e:
-        logger.error(f"生成简易配置注释错误...", e=e)
+        logger.error("生成简易配置注释错误...", e=e)
     if temp_file.exists():
         temp_file.unlink()
 

zhenxun/builtin_plugins/init/init_plugin.py

@@ -21,6 +21,8 @@ from zhenxun.utils.enum import (
     PluginLimitType,
 )
 
+from .manager import manager
+
 _yaml = YAML(pure=True)
 _yaml.allow_unicode = True
 _yaml.indent = 2
@@ -148,23 +150,25 @@ async def _():
     #     ["name", "author", "version", "admin_level", "plugin_type"],
     #     10,
     # )
-    if limit_list:
-        limit_create = []
-        plugins = []
-        if module_path_list := [limit.module_path for limit in limit_list]:
-            plugins = await PluginInfo.filter(module_path__in=module_path_list).all()
-        if plugins:
-            for limit in limit_list:
-                if lmt := [p for p in plugins if p.module_path == limit.module_path]:
-                    plugin = lmt[0]
-                    limit_type_list = [
-                        _limit.limit_type for _limit in await plugin.plugin_limit.all()  # type: ignore
-                    ]
-                    if limit.limit_type not in limit_type_list:
-                        limit.plugin = plugin
-                        limit_create.append(limit)
-        if limit_create:
-            await PluginLimit.bulk_create(limit_create, 10)
+    # for limit in limit_list:
+    #     limit_create = []
+    #     plugins = []
+    #     if module_path_list := [limit.module_path for limit in limit_list]:
+    #         plugins = await PluginInfo.get_plugins(module_path__in=module_path_list)
+    #     if plugins:
+    #         for limit in limit_list:
+    #             if lmt := [p for p in plugins if p.module_path == limit.module_path]:
+    #                 plugin = lmt[0]
+    #                 """不在数据库中"""
+    #                 limit_type_list = [
+    #                     _limit.limit_type
+    #                     for _limit in await plugin.plugin_limit.all()  # type: ignore
+    #                 ]
+    #                 if limit.limit_type not in limit_type_list:
+    #                     limit.plugin = plugin
+    #                     limit_create.append(limit)
+    #     if limit_create:
+    #         await PluginLimit.bulk_create(limit_create, 10)
     if task_list:
         module_dict = {
             t[1]: t[0] for t in await TaskInfo.all().values_list("id", "module")
@@ -195,10 +199,18 @@ async def _():
     await data_migration()
     await PluginInfo.filter(module_path__in=load_plugin).update(load_status=True)
     await PluginInfo.filter(module_path__not_in=load_plugin).update(load_status=False)
+    manager.init()
+    if limit_list:
+        for limit in limit_list:
+            if not manager.exist(limit.module_path, limit.limit_type):
+                """不存在，添加"""
+                manager.add(limit.module_path, limit)
+    manager.save_file()
+    await manager.load_to_db()
 
 
 async def data_migration():
-    await limit_migration()
+    # await limit_migration()
     await plugin_migration()
     await group_migration()
 

zhenxun/builtin_plugins/init/manager.py

@@ -0,0 +1,417 @@
+from copy import deepcopy
+
+from ruamel.yaml import YAML
+
+from zhenxun.services.log import logger
+from zhenxun.configs.path_config import DATA_PATH
+from zhenxun.models.plugin_info import PluginInfo
+from zhenxun.models.plugin_limit import PluginLimit
+from zhenxun.utils.enum import BlockType, LimitCheckType, PluginLimitType
+from zhenxun.configs.utils import BaseBlock, PluginCdBlock, PluginCountBlock
+
+_yaml = YAML(pure=True)
+_yaml.indent = 2
+_yaml.allow_unicode = True
+
+
+CD_TEST = """需要cd的功能
+自定义的功能需要cd也可以在此配置
+key：模块名称
+cd：cd 时长（秒）
+status：此限制的开关状态
+check_type：'PRIVATE'/'GROUP'/'ALL'，限制私聊/群聊/全部
+watch_type：监听对象，以user_id或group_id作为键来限制，'USER'：用户id，'GROUP'：群id
+                                 示例：'USER':用户N秒内触发1次，'GROUP':群N秒内触发1次
+result：回复的话,可以添加[at],[uname],[nickname]来对应艾特，用户群名称，昵称系统昵称
+result 为 "" 或 None 时则不回复
+result示例："[uname]你冲的太快了，[nickname]先生，请稍后再冲[at]"
+result回复："老色批你冲的太快了，欧尼酱先生，请稍后再冲@老色批"
+     用户昵称↑     昵称系统的昵称↑          艾特用户↑"""
+
+
+BLOCK_TEST = """用户调用阻塞
+即 当用户调用此功能还未结束时
+用发送消息阻止用户重复调用此命令直到该命令结束
+key：模块名称
+status：此限制的开关状态
+check_type：'PRIVATE'/'GROUP'/'ALL'，限制私聊/群聊/全部
+watch_type：监听对象，以user_id或group_id作为键来限制，'USER'：用户id，'GROUP'：群id
+                                    示例：'USER'：阻塞用户，'group'：阻塞群聊
+result：回复的话，可以添加[at]，[uname]，[nickname]来对应艾特，用户群名称，昵称系统昵称
+result 为 "" 或 None 时则不回复
+result示例："[uname]你冲的太快了，[nickname]先生，请稍后再冲[at]"
+result回复："老色批你冲的太快了，欧尼酱先生，请稍后再冲@老色批"
+     用户昵称↑     昵称系统的昵称↑          艾特用户↑"""
+
+COUNT_TEST = """命令每日次数限制
+即 用户/群聊 每日可调用命令的次数 [数据内存存储，重启将会重置]
+每日调用直到 00:00 刷新
+key：模块名称
+max_count: 每日调用上限
+status：此限制的开关状态
+watch_type：监听对象，以user_id或group_id作为键来限制，'USER'：用户id，'GROUP'：群id
+                                     示例：'USER'：用户上限，'group'：群聊上限
+result：回复的话，可以添加[at]，[uname]，[nickname]来对应艾特，用户群名称，昵称系统昵称
+result 为 "" 或 None 时则不回复
+result示例："[uname]你冲的太快了，[nickname]先生，请稍后再冲[at]"
+result回复："老色批你冲的太快了，欧尼酱先生，请稍后再冲@老色批"
+     用户昵称↑     昵称系统的昵称↑          艾特用户↑"""
+
+
+class Manager:
+    """
+    插件命令 cd 管理器
+    """
+
+    def __init__(self):
+        self.cd_file = DATA_PATH / "configs" / "plugins2cd.yaml"
+        self.block_file = DATA_PATH / "configs" / "plugins2block.yaml"
+        self.count_file = DATA_PATH / "configs" / "plugins2count.yaml"
+        self.cd_data = {}
+        self.block_data = {}
+        self.count_data = {}
+
+    def add(
+        self,
+        module_path: str,
+        data: BaseBlock | PluginCdBlock | PluginCountBlock | PluginLimit,
+    ):
+        """添加限制"""
+        if isinstance(data, PluginLimit):
+            check_type = BlockType.ALL
+            if LimitCheckType.GROUP == data.check_type:
+                check_type = BlockType.GROUP
+            elif LimitCheckType.PRIVATE == data.check_type:
+                check_type = BlockType.PRIVATE
+            if data.limit_type == PluginLimitType.CD:
+                data = PluginCdBlock(
+                    status=data.status,
+                    check_type=check_type,
+                    watch_type=data.watch_type,
+                    result=data.result,
+                    cd=data.cd,
+                )
+            elif data.limit_type == PluginLimitType.BLOCK:
+                data = BaseBlock(
+                    status=data.status,
+                    check_type=check_type,
+                    watch_type=data.watch_type,
+                    result=data.result,
+                )
+            elif data.limit_type == PluginLimitType.COUNT:
+                data = PluginCountBlock(
+                    status=data.status,
+                    watch_type=data.watch_type,
+                    result=data.result,
+                    max_count=data.max_count,
+                )
+        if isinstance(data, PluginCdBlock):
+            self.cd_data[module_path] = data
+        elif isinstance(data, PluginCountBlock):
+            self.count_data[module_path] = data
+        elif isinstance(data, BaseBlock):
+            self.block_data[module_path] = data
+
+    def exist(self, module_path: str, type: PluginLimitType):
+        """是否存在"""
+        if type == PluginLimitType.CD:
+            return module_path in self.cd_data
+        elif type == PluginLimitType.BLOCK:
+            return module_path in self.block_data
+        elif type == PluginLimitType.COUNT:
+            return module_path in self.count_data
+
+    def init(self):
+        if not self.cd_file.exists():
+            self.save_cd_file()
+        if not self.block_file.exists():
+            self.save_block_file()
+        if not self.count_file.exists():
+            self.save_count_file()
+        self.__load_file()
+
+    def __load_file(self):
+        self.__load_block_file()
+        self.__load_cd_file()
+        self.__load_count_file()
+
+    def save_file(self):
+        """保存文件"""
+        self.save_cd_file()
+        self.save_block_file()
+        self.save_count_file()
+
+    def save_cd_file(self):
+        """保存文件"""
+        self._extracted_from_save_file_3("PluginCdLimit", CD_TEST, self.cd_data)
+
+    def save_block_file(self):
+        """保存文件"""
+        self._extracted_from_save_file_3(
+            "PluginBlockLimit", BLOCK_TEST, self.block_data
+        )
+
+    def save_count_file(self):
+        """保存文件"""
+        self._extracted_from_save_file_3(
+            "PluginCountLimit", COUNT_TEST, self.count_data
+        )
+
+    def _extracted_from_save_file_3(self, type_: str, after: str, data: dict):
+        """保存文件
+
+        参数:
+            type_: 类型参数
+            after: 备注
+        """
+        temp_data = deepcopy(data)
+        if not temp_data:
+            temp_data = {
+                "test": {
+                    "status": False,
+                    "check_type": "ALL",
+                    "limit_type": "USER",
+                    "result": "你冲的太快了，请稍后再冲",
+                }
+            }
+            if type_ == "PluginCdLimit":
+                temp_data["test"]["cd"] = 5
+            elif type_ == "PluginCountLimit":
+                temp_data["test"]["max_count"] = 5
+                del temp_data["test"]["check_type"]
+        else:
+            for v in temp_data:
+                temp_data[v] = temp_data[v].dict()
+                if check_type := temp_data[v].get("check_type"):
+                    temp_data[v]["check_type"] = str(check_type)
+                if watch_type := temp_data[v].get("watch_type"):
+                    temp_data[v]["watch_type"] = str(watch_type)
+                if type_ == "PluginCountLimit":
+                    del temp_data[v]["check_type"]
+        file = self.block_file
+        if type_ == "PluginCdLimit":
+            file = self.cd_file
+        elif type_ == "PluginCountLimit":
+            file = self.count_file
+        with open(file, "w", encoding="utf8") as f:
+            _yaml.dump({type_: temp_data}, f)
+        with open(file, encoding="utf8") as rf:
+            _data = _yaml.load(rf)
+        _data.yaml_set_comment_before_after_key(after=after, key=type_)
+        with open(file, "w", encoding="utf8") as wf:
+            _yaml.dump(_data, wf)
+
+    def __load_cd_file(self):
+        self.cd_data: dict[str, PluginCdBlock] = {}
+        if self.cd_file.exists():
+            with open(self.cd_file, encoding="utf8") as f:
+                temp = _yaml.load(f)
+                if "PluginCdLimit" in temp.keys():
+                    for k, v in temp["PluginCdLimit"].items():
+                        self.cd_data[k] = PluginCdBlock.parse_obj(v)
+
+    def __load_block_file(self):
+        self.block_data: dict[str, BaseBlock] = {}
+        if self.block_file.exists():
+            with open(self.block_file, encoding="utf8") as f:
+                temp = _yaml.load(f)
+                if "PluginBlockLimit" in temp.keys():
+                    for k, v in temp["PluginBlockLimit"].items():
+                        self.block_data[k] = BaseBlock.parse_obj(v)
+
+    def __load_count_file(self):
+        self.count_data: dict[str, PluginCountBlock] = {}
+        if self.count_file.exists():
+            with open(self.count_file, encoding="utf8") as f:
+                temp = _yaml.load(f)
+                if "PluginCountLimit" in temp.keys():
+                    for k, v in temp["PluginCountLimit"].items():
+                        self.count_data[k] = PluginCountBlock.parse_obj(v)
+
+    def __replace_data(
+        self,
+        db_data: PluginLimit | None,
+        limit: PluginCdBlock | BaseBlock | PluginCountBlock,
+    ) -> PluginLimit:
+        """替换数据"""
+        if not db_data:
+            db_data = PluginLimit()
+        db_data.status = limit.status
+        check_type = LimitCheckType.ALL
+        if BlockType.GROUP == limit.check_type:
+            check_type = LimitCheckType.GROUP
+        elif BlockType.PRIVATE == limit.check_type:
+            check_type = LimitCheckType.PRIVATE
+        db_data.check_type = check_type
+        db_data.watch_type = limit.watch_type
+        db_data.result = limit.result or ""
+        return db_data
+
+    def __set_data(
+        self,
+        k: str,
+        db_data: PluginLimit | None,
+        limit: PluginCdBlock | BaseBlock | PluginCountBlock,
+        limit_type: PluginLimitType,
+        module2plugin: dict[str, PluginInfo],
+    ) -> tuple[PluginLimit, bool]:
+        """设置数据
+
+        参数:
+            k: 模块名
+            db_data: 数据库数据
+            limit: 文件数据
+            limit_type: 限制类型
+            module2plugin: 模块:插件信息
+
+        返回:
+            tuple[PluginLimit, bool]: PluginLimit，是否创建
+        """
+        if not db_data:
+            return (
+                PluginLimit(
+                    module=k.split(".")[-1],
+                    module_path=k,
+                    limit_type=limit_type,
+                    plugin=module2plugin.get(k),
+                    cd=getattr(limit, "cd", None),
+                    max_count=getattr(limit, "max_count", None),
+                    status=limit.status,
+                    check_type=limit.check_type,
+                    watch_type=limit.watch_type,
+                    result=limit.result,
+                ),
+                True,
+            )
+        db_data = self.__replace_data(db_data, limit)
+        if limit_type == PluginLimitType.CD:
+            db_data.cd = limit.cd  # type: ignore
+        if limit_type == PluginLimitType.COUNT:
+            db_data.max_count = limit.max_count  # type: ignore
+        return db_data, False
+
+    def __get_file_data(self, limit_type: PluginLimitType) -> dict:
+        """获取文件数据
+
+        参数:
+            limit_type: 限制类型
+
+        返回:
+            dict: 文件数据
+        """
+        if limit_type == PluginLimitType.CD:
+            return self.cd_data
+        elif limit_type == PluginLimitType.COUNT:
+            return self.count_data
+        else:
+            return self.block_data
+
+    def __set_db_limits(
+        self,
+        db_limits: list[PluginLimit],
+        module2plugin: dict[str, PluginInfo],
+        limit_type: PluginLimitType,
+    ) -> tuple[list[PluginLimit], list[PluginLimit], list[int]]:
+        """更新cd限制数据
+
+        参数:
+            db_limits: 数据库limits
+            module2plugin: 模块:插件信息
+
+        返回:
+            tuple[list[PluginLimit], list[PluginLimit]]: 创建列表，更新列表
+        """
+        update_list = []
+        create_list = []
+        delete_list = []
+        db_type_limits = [
+            limit for limit in db_limits if limit.limit_type == limit_type
+        ]
+        if data := self.__get_file_data(limit_type):
+            db_type_limit_modules = [
+                (limit.module_path, limit.id) for limit in db_type_limits
+            ]
+            delete_list.extend(
+                id
+                for module_path, id in db_type_limit_modules
+                if module_path not in data.keys()
+            )
+            for k, v in data.items():
+                if not module2plugin.get(k):
+                    if k != "test":
+                        logger.warning(
+                            f"插件模块 {k} 未加载，已过滤当前 {v._type} 限制..."
+                        )
+                    continue
+                db_data = [limit for limit in db_type_limits if limit.module_path == k]
+                db_data, is_create = self.__set_data(
+                    k, db_data[0] if db_data else None, v, limit_type, module2plugin
+                )
+                if is_create:
+                    create_list.append(db_data)
+                else:
+                    update_list.append(db_data)
+        else:
+            delete_list = [limit.id for limit in db_type_limits]
+        return create_list, update_list, delete_list
+
+    async def __set_all_limit(
+        self,
+    ) -> tuple[list[PluginLimit], list[PluginLimit], list[int]]:
+        """获取所有插件限制数据
+
+        返回:
+            tuple[list[PluginLimit], list[PluginLimit]]: 创建列表，更新列表
+        """
+        db_limits = await PluginLimit.all()
+        modules = set(
+            list(self.cd_data.keys())
+            + list(self.block_data.keys())
+            + list(self.count_data.keys())
+        )
+        plugins = await PluginInfo.get_plugins(module_path__in=modules)
+        module2plugin = {p.module_path: p for p in plugins}
+        create_list, update_list, delete_list = self.__set_db_limits(
+            db_limits, module2plugin, PluginLimitType.CD
+        )
+        create_list1, update_list1, delete_list1 = self.__set_db_limits(
+            db_limits, module2plugin, PluginLimitType.COUNT
+        )
+        create_list2, update_list2, delete_list2 = self.__set_db_limits(
+            db_limits, module2plugin, PluginLimitType.BLOCK
+        )
+        all_create = create_list + create_list1 + create_list2
+        all_update = update_list + update_list1 + update_list2
+        all_delete = delete_list + delete_list1 + delete_list2
+        return all_create, all_update, all_delete
+
+    async def load_to_db(self):
+        """读取配置文件"""
+
+        create_list, update_list, delete_list = await self.__set_all_limit()
+        if create_list:
+            await PluginLimit.bulk_create(create_list)
+        if update_list:
+            for limit in update_list:
+                await limit.save(
+                    update_fields=[
+                        "status",
+                        "check_type",
+                        "watch_type",
+                        "result",
+                        "cd",
+                        "max_count",
+                    ]
+                )
+            # TODO: tortoise.exceptions.OperationalError:syntax error at or near "GROUP"
+            # await PluginLimit.bulk_update(
+            #     update_list,
+            #     ["status", "check_type", "watch_type", "result", "cd", "max_count"],
+            # )
+        if delete_list:
+            await PluginLimit.filter(id__in=delete_list).delete()
+        cnt = await PluginLimit.filter(status=True).count()
+        logger.info(f"已经加载 {cnt} 个插件限制.")
+
+
+manager = Manager()

zhenxun/configs/utils/__init__.py

@@ -408,13 +408,6 @@ class ConfigsManager:
         """
         if save_simple_data:
             with open(self._simple_file, "w", encoding="utf8") as f:
-                # yaml.dump(
-                #     self._simple_data,
-                #     f,
-                #     indent=2,
-                #     Dumper=yaml.RoundTripDumper,
-                #     allow_unicode=True,
-                # )
                 _yaml.dump(self._simple_data, f)
         path = path or self.file
         data = {}
@@ -426,14 +419,10 @@ class ConfigsManager:
                 del value["arg_parser"]
                 data[module][config] = value
         with open(path, "w", encoding="utf8") as f:
-            # yaml.dump(
-            #     data, f, indent=2, Dumper=yaml.RoundTripDumper, allow_unicode=True
-            # )
             _yaml.dump(data, f)
 
     def reload(self):
         """重新加载配置文件"""
-        _yaml = YAML()
         if self._simple_file.exists():
             with open(self._simple_file, encoding="utf8") as f:
                 self._simple_data = _yaml.load(f)

zhenxun/models/plugin_info.py

@@ -1,3 +1,5 @@
+from typing_extensions import Self
+
 from tortoise import fields
 
 from zhenxun.services.db_context import Model
@@ -51,6 +53,30 @@ class PluginInfo(Model):
         table = "plugin_info"
         table_description = "插件基本信息"
 
+    @classmethod
+    async def get_plugin(cls, load_status: bool = True, **kwargs) -> Self | None:
+        """获取插件列表
+
+        参数:
+            load_status: 加载状态.
+
+        返回:
+            Self | None: 插件
+        """
+        return await cls.get_or_none(load_status=load_status, **kwargs)
+
+    @classmethod
+    async def get_plugins(cls, load_status: bool = True, **kwargs) -> list[Self]:
+        """获取插件列表
+
+        参数:
+            load_status: 加载状态.
+
+        返回:
+            list[Self]: 插件列表
+        """
+        return await cls.filter(load_status=load_status, **kwargs).all()
+
     @classmethod
     async def _run_script(cls):
         return [

zhenxun/models/plugin_limit.py

@@ -35,6 +35,6 @@ class PluginLimit(Model):
     max_count = fields.IntField(null=True, description="最大调用次数")
     """最大调用次数"""
 
-    class Meta:
+    class Meta:  # type: ignore
         table = "plugin_limit"
         table_description = "插件限制"