zhenxun_bot/zhenxun/builtin_plugins/init/init_plugin.py

import aiofiles
import nonebot
from nonebot import get_loaded_plugins
from nonebot.drivers import Driver
from nonebot.plugin import Plugin, PluginMetadata
from ruamel.yaml import YAML
import ujson as json

from zhenxun.configs.path_config import DATA_PATH
from zhenxun.configs.utils import PluginExtraData, PluginSetting
from zhenxun.models.group_console import GroupConsole
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.models.plugin_limit import PluginLimit
from zhenxun.models.task_info import TaskInfo
from zhenxun.services.log import logger
from zhenxun.utils.enum import (
    BlockType,
    LimitCheckType,
    LimitWatchType,
    PluginLimitType,
    PluginType,
)

from .manager import manager

_yaml = YAML(pure=True)
_yaml.allow_unicode = True
_yaml.indent = 2

driver: Driver = nonebot.get_driver()


async def _handle_setting(
    plugin: Plugin,
    plugin_list: list[PluginInfo],
    limit_list: list[PluginLimit],
    task_list: list[tuple[bool, TaskInfo]],
):
    """处理插件设置

    参数:
        plugin: Plugin
        plugin_list: 插件列表
        limit_list: 插件限制列表
    """
    metadata = plugin.metadata
    if not metadata:
        if not plugin.sub_plugins:
            return
        """父插件"""
        metadata = PluginMetadata(name=plugin.name, description="", usage="")
    extra = metadata.extra
    extra_data = PluginExtraData(**extra)
    logger.debug(f"{metadata.name}:{plugin.name} -> {extra}", "初始化插件数据")
    setting = extra_data.setting or PluginSetting()
    if metadata.type == "library":
        extra_data.plugin_type = PluginType.HIDDEN
    if extra_data.plugin_type == PluginType.HIDDEN:
        extra_data.menu_type = ""
    if plugin.sub_plugins:
        extra_data.plugin_type = PluginType.PARENT
    plugin_list.append(
        PluginInfo(
            module=plugin.name,
            module_path=plugin.module_name,
            name=metadata.name,
            author=extra_data.author,
            version=extra_data.version,
            level=setting.level,
            default_status=setting.default_status,
            limit_superuser=setting.limit_superuser,
            menu_type=extra_data.menu_type,
            cost_gold=setting.cost_gold,
            plugin_type=extra_data.plugin_type,
            admin_level=extra_data.admin_level,
            parent=(plugin.parent_plugin.module_name if plugin.parent_plugin else None),
        )
    )
    if extra_data.limits:
        limit_list.extend(
            PluginLimit(
                module=plugin.name,
                module_path=plugin.module_name,
                limit_type=limit._type,
                watch_type=limit.watch_type,
                status=limit.status,
                check_type=limit.check_type,
                result=limit.result,
                cd=getattr(limit, "cd", None),
                max_count=getattr(limit, "max_count", None),
            )
            for limit in extra_data.limits
        )
    if extra_data.tasks:
        task_list.extend(
            (
                task.create_status,
                TaskInfo(
                    module=task.module,
                    name=task.name,
                    status=task.status,
                    run_time=task.run_time,
                    default_status=task.default_status,
                ),
            )
            for task in extra_data.tasks
        )


@driver.on_startup
async def _():
    """
    初始化插件数据配置
    """
    plugin_list: list[PluginInfo] = []
    limit_list: list[PluginLimit] = []
    task_list = []
    module2id = {}
    load_plugin = []
    if module_list := await PluginInfo.all().values("id", "module_path"):
        module2id = {m["module_path"]: m["id"] for m in module_list}
    for plugin in get_loaded_plugins():
        load_plugin.append(plugin.module_name)
        await _handle_setting(plugin, plugin_list, limit_list, task_list)
    create_list = []
    update_list = []
    for plugin in plugin_list:
        if plugin.module_path not in module2id:
            create_list.append(plugin)
        else:
            plugin.id = module2id[plugin.module_path]
            await plugin.save(
                update_fields=[
                    "name",
                    "author",
                    "version",
                    "admin_level",
                    "plugin_type",
                    "is_show",
                ]
            )
            update_list.append(plugin)
    if create_list:
        await PluginInfo.bulk_create(create_list, 10)
    # if update_list:
    #     # TODO: 批量更新无法更新plugin_type: tortoise.exceptions.OperationalError:
    #           column "superuser" does not exist
    #     pass
    # await PluginInfo.bulk_update(
    #     update_list,
    #     ["name", "author", "version", "admin_level", "plugin_type"],
    #     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")
        }
        create_list = []
        update_list = []
        for status, task in task_list:
            if task.module not in module_dict:
                create_list.append((status, task))
            else:
                task.id = module_dict[task.module]
                update_list.append(task)
        if create_list:
            _create_list = [t[1] for t in create_list]
            await TaskInfo.bulk_create(_create_list, 10)
            if block := [t[1].module for t in create_list if not t[0]]:
                block_task = ",".join(block) + ","
                if group_list := await GroupConsole.all():
                    for group in group_list:
                        group.block_task += block_task
                    await GroupConsole.bulk_update(group_list, ["block_task"], 10)
        if update_list:
            await TaskInfo.bulk_update(
                update_list,
                ["run_time", "name"],
                10,
            )
    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.exists(limit.module, limit.limit_type):
                """不存在，添加"""
                manager.add(limit.module, limit)
    manager.save_file()
    await manager.load_to_db()


async def data_migration():
    # await limit_migration()
    await plugin_migration()
    await group_migration()


async def limit_migration():
    """插件限制迁移"""
    cd_file = DATA_PATH / "configs" / "plugins2cd.yaml"
    block_file = DATA_PATH / "configs" / "plugins2block.yaml"
    count_file = DATA_PATH / "configs" / "plugins2count.yaml"
    limit_data: dict[str, list[tuple[str, dict]]] = {}
    if cd_file.exists():
        async with aiofiles.open(cd_file, encoding="utf8") as f:
            if data := _yaml.load(await f.read()):
                for k in data["PluginCdLimit"]:
                    limit_data[k] = [("CD", data["PluginCdLimit"][k])]
        cd_file.unlink()
    if block_file.exists():
        async with aiofiles.open(block_file, encoding="utf8") as f:
            if data := _yaml.load(await f.read()):
                for k in data["PluginBlockLimit"]:
                    if k in limit_data:
                        limit_data[k].append(("BLOCK", data["PluginBlockLimit"][k]))
                    else:
                        limit_data[k] = [("BLOCK", data["PluginBlockLimit"][k])]
        block_file.unlink()
    if count_file.exists():
        async with aiofiles.open(count_file, encoding="utf8") as f:
            if data := _yaml.load(await f.read()):
                for k in data["PluginCountLimit"]:
                    if k in limit_data:
                        limit_data[k].append(("COUNT", data["PluginCountLimit"][k]))
                    else:
                        limit_data[k] = [("COUNT", data["PluginCountLimit"][k])]
        count_file.unlink()
    if limit_data:
        logger.info("开始迁移插件限制数据...")
        update_list = []
        create_list = []
        plugins = await PluginInfo.filter(module__in=limit_data.keys())
        for plugin in plugins:
            limits: list[PluginLimit] = await plugin.plugin_limit.all()  # type: ignore
            exits_limit = [x[0] for x in limit_data[plugin.module]]
            _not_create_type = []
            for limit in limits:
                if _limit_list := [
                    x[1]
                    for x in limit_data[plugin.module]
                    if x[0] == str(limit.limit_type)
                ]:
                    """修改"""
                    _not_create_type.append(str(limit.limit_type))
                    _limit = _limit_list[0]
                    watch_type = LimitWatchType.USER
                    if _limit.get("watch_type") == "group":
                        watch_type = LimitWatchType.GROUP
                    check_type = LimitCheckType.ALL
                    if _limit.get("check_type") == "private":
                        check_type = LimitCheckType.PRIVATE
                    elif _limit.get("check_type") == "group":
                        check_type = LimitCheckType.GROUP
                    limit.watch_type = watch_type
                    limit.result = _limit.get("rst", "")
                    limit.status = _limit.get("status", True)
                    if limit.watch_type != PluginLimitType.COUNT:
                        limit.check_type = check_type
                    if limit.watch_type == PluginLimitType.CD:
                        limit.cd = _limit["cd"]
                    if limit.watch_type == PluginLimitType.COUNT:
                        limit.max_count = _limit["count"]
                    await limit.save()
                    update_list.append(limit)
            for s in [e for e in exits_limit if e not in _not_create_type]:
                if _limit_list := [
                    x[1] for x in limit_data[plugin.module] if s == x[0]
                ]:
                    _limit = _limit_list[0]
                    limit_type = PluginLimitType.CD
                    if s == "BLOCK":
                        limit_type = PluginLimitType.BLOCK
                    elif s == "COUNT":
                        limit_type = PluginLimitType.COUNT
                    watch_type = LimitWatchType.USER
                    if _limit.get("watch_type") == "group":
                        watch_type = LimitWatchType.GROUP
                    check_type = LimitCheckType.ALL
                    if _limit.get("check_type") == "private":
                        check_type = LimitCheckType.PRIVATE
                    elif _limit.get("check_type") == "group":
                        check_type = LimitCheckType.GROUP
                    create_list.append(
                        PluginLimit(
                            module=plugin.module,
                            module_path=plugin.module_path,
                            plugin=plugin,
                            limit_type=limit_type,
                            watch_type=watch_type,
                            status=_limit.get("status", True),
                            check_type=check_type,
                            result=_limit.get("rst", ""),
                            cd=_limit.get("cd"),
                            max_count=_limit.get("max_count"),
                        )
                    )
        # TODO: 批量错误 tortoise.exceptions.OperationalError:
        # syntax error at or near "ALL"
        # if update_list:
        #     await PluginLimit.bulk_update(
        #         update_list,
        #         [
        #             "watch_type",
        #             "status",
        #             "check_type",
        #             "result",
        #             "cd",
        #             "max_count",
        #         ],
        #         10,
        #     )
        if create_list:
            await PluginLimit.bulk_create(create_list, 10)
        logger.info("迁移插件限制数据完成!")


async def plugin_migration():
    """迁移插件数据"""
    setting_file = DATA_PATH / "configs" / "plugins2settings.yaml"
    plugin_file = DATA_PATH / "manager" / "plugins_manager.json"
    if setting_file.exists():
        async with aiofiles.open(setting_file, encoding="utf8") as f:
            if data := _yaml.load(await f.read()):
                logger.info("开始迁移插件setting数据...")
                data = data["PluginSettings"]
                plugins = await PluginInfo.filter(module__in=data.keys())
                for plugin in plugins:
                    if plugin_data_list := [
                        data[p] for p in data if p == plugin.module
                    ]:
                        plugin_data = plugin_data_list[0]
                        plugin.default_status = plugin_data.get("default_status", True)
                        plugin.level = plugin_data.get("level", 5)
                        plugin.limit_superuser = plugin_data.get(
                            "limit_superuser", False
                        )
                        plugin.menu_type = plugin_data.get("plugin_type", ["功能"])[0]
                        plugin.cost_gold = plugin_data.get("cost_gold", 0)
                await PluginInfo.bulk_update(
                    plugins,
                    [
                        "default_status",
                        "level",
                        "limit_superuser",
                        "menu_type",
                        "cost_gold",
                    ],
                    10,
                )
        setting_file.unlink()
        logger.info("迁移插件setting数据完成!")
    if plugin_file.exists():
        async with aiofiles.open(plugin_file, encoding="utf8") as f:
            if data := json.loads(await f.read()):
                logger.info("开始迁移插件数据...")
                plugins = await PluginInfo.filter(module__in=data.keys())
                for plugin in plugins:
                    if plugin_data := data.get(plugin.module):
                        plugin.status = plugin_data.get("status", True)
                        block_type = None
                        get_block = plugin_data.get("block_type")
                        if get_block == "all":
                            block_type = BlockType.ALL
                        elif get_block == "private":
                            block_type = BlockType.PRIVATE
                        elif get_block == "group":
                            block_type = BlockType.GROUP
                        plugin.block_type = block_type
                        await plugin.save(update_fields=["status", "block_type"])
                # TODO: tortoise.exceptions.OperationalError: syntax error at
                # or near "ALL"
                # await PluginInfo.bulk_update(plugins, ["status", "block_type"], 10)
        plugin_file.unlink()
        logger.info("迁移插件数据完成!")


async def group_migration():
    """
    群组数据迁移
    """
    group_file = DATA_PATH / "manager" / "group_manager.json"
    if group_file.exists():
        async with aiofiles.open(group_file, encoding="utf8") as f:
            if data := json.loads(await f.read()):
                logger.info("开始迁移群组数据...")
                update_list = []
                create_list = []
                white_group = data["white_group"]
                old_group_list: dict = data["group_manager"]
                if close_task := data["close_task"]:
                    """全局被动关闭"""
                    await TaskInfo.filter(module__in=close_task).update(status=False)
                group_list = await GroupConsole.filter(
                    group_id__in=old_group_list.keys()
                )
                for old_group_id, old_group in old_group_list.items():
                    block_plugin = ""
                    block_task = ""
                    status = old_group.get("status", True)
                    level = old_group.get("level", 5)
                    if close_plugins := old_group.get("close_plugins"):
                        block_plugin = ",".join(close_plugins) + ","
                    if group_task_status := old_group.get("group_task_status"):
                        close_task = [
                            t for t in group_task_status if not group_task_status[t]
                        ]
                        block_task = ",".join(close_task) + ","
                    if group_ := [g for g in group_list if g.group_id == old_group_id]:
                        group = group_[0]
                        if group.group_id in white_group:
                            group.is_super = True
                        group.status = status
                        group.block_plugin = block_plugin
                        group.block_task = block_task
                        group.level = level
                        update_list.append(group)
                    else:
                        """添加"""
                        create_list.append(
                            GroupConsole(
                                group_id=old_group_id,
                                status=status,
                                level=level,
                                block_plugin=block_plugin,
                                block_task=block_task,
                                is_super=old_group_id in white_group,
                            )
                        )
                if update_list:
                    await GroupConsole.bulk_update(
                        update_list,
                        ["is_super", "status", "block_plugin", "block_task"],
                        10,
                    )
                if create_list:
                    await GroupConsole.bulk_create(create_list, 10)
        group_file.unlink()
        logger.info("迁移群组数据完成!")