Diuse/zhenxun_bot
1
0
Fork 0
mirror of https://github.com/zhenxun-org/zhenxun_bot.git synced 2025-12-15 14:22:55 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

✨ 构建缓存，hook使用缓存

Browse Source
This commit is contained in:
HibiKier 2025-01-08 15:23:10 +08:00
parent 45649bb29d
commit 521bcaceeb
16 changed files with 996 additions and 660 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

580
zhenxun/builtin_plugins/hooks/_auth_checker.py
View File

@ -1,580 +0,0 @@
from nonebot.adapters import Bot, Event
from nonebot.adapters.onebot.v11 import PokeNotifyEvent
from nonebot.exception import IgnoredException
from nonebot.matcher import Matcher
from nonebot_plugin_alconna import At, UniMsg
from nonebot_plugin_session import EventSession
from pydantic import BaseModel
from tortoise.exceptions import IntegrityError
from zhenxun.configs.config import Config
from zhenxun.models.bot_console import BotConsole
from zhenxun.models.group_console import GroupConsole
from zhenxun.models.level_user import LevelUser
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.models.plugin_limit import PluginLimit
from zhenxun.models.user_console import UserConsole
from zhenxun.services.log import logger
from zhenxun.utils.cache_utils import Cache
from zhenxun.utils.enum import (
BlockType,
CacheType,
GoldHandle,
LimitWatchType,
PluginLimitType,
PluginType,
)
from zhenxun.utils.exception import InsufficientGold
from zhenxun.utils.message import MessageUtils
from zhenxun.utils.utils import CountLimiter, FreqLimiter, UserBlockLimiter
base_config = Config.get("hook")
class Limit(BaseModel):
limit: PluginLimit
limiter: FreqLimiter | UserBlockLimiter | CountLimiter
class Config:
arbitrary_types_allowed = True
class LimitManage:
add_module = [] # noqa: RUF012
cd_limit: dict[str, Limit] = {} # noqa: RUF012
block_limit: dict[str, Limit] = {} # noqa: RUF012
count_limit: dict[str, Limit] = {} # noqa: RUF012
@classmethod
def add_limit(cls, limit: PluginLimit):
"""添加限制
参数:
limit: PluginLimit
"""
if limit.module not in cls.add_module:
cls.add_module.append(limit.module)
if limit.limit_type == PluginLimitType.BLOCK:
cls.block_limit[limit.module] = Limit(
limit=limit, limiter=UserBlockLimiter()
)
elif limit.limit_type == PluginLimitType.CD:
cls.cd_limit[limit.module] = Limit(
limit=limit, limiter=FreqLimiter(limit.cd)
)
elif limit.limit_type == PluginLimitType.COUNT:
cls.count_limit[limit.module] = Limit(
limit=limit, limiter=CountLimiter(limit.max_count)
)
@classmethod
def unblock(
cls, module: str, user_id: str, group_id: str | None, channel_id: str | None
):
"""解除插件block
参数:
module: 模块名
user_id: 用户id
group_id: 群组id
channel_id: 频道id
"""
if limit_model := cls.block_limit.get(module):
limit = limit_model.limit
limiter: UserBlockLimiter = limit_model.limiter # type: ignore
key_type = user_id
if group_id and limit.watch_type == LimitWatchType.GROUP:
key_type = channel_id or group_id
logger.debug(
f"解除对象: {key_type} 的block限制",
"AuthChecker",
session=user_id,
group_id=group_id,
)
limiter.set_false(key_type)
@classmethod
async def check(
cls,
module: str,
user_id: str,
group_id: str | None,
channel_id: str | None,
session: EventSession,
):
"""检测限制
参数:
module: 模块名
user_id: 用户id
group_id: 群组id
channel_id: 频道id
session: Session
异常:
IgnoredException: IgnoredException
"""
if limit_model := cls.cd_limit.get(module):
await cls.__check(limit_model, user_id, group_id, channel_id, session)
if limit_model := cls.block_limit.get(module):
await cls.__check(limit_model, user_id, group_id, channel_id, session)
if limit_model := cls.count_limit.get(module):
await cls.__check(limit_model, user_id, group_id, channel_id, session)
@classmethod
async def __check(
cls,
limit_model: Limit | None,
user_id: str,
group_id: str | None,
channel_id: str | None,
session: EventSession,
):
"""检测限制
参数:
limit_model: Limit
user_id: 用户id
group_id: 群组id
channel_id: 频道id
session: Session
异常:
IgnoredException: IgnoredException
"""
if not limit_model:
return
limit = limit_model.limit
limiter = limit_model.limiter
is_limit = (
LimitWatchType.ALL
or (group_id and limit.watch_type == LimitWatchType.GROUP)
or (not group_id and limit.watch_type == LimitWatchType.USER)
)
key_type = user_id
if group_id and limit.watch_type == LimitWatchType.GROUP:
key_type = channel_id or group_id
if is_limit and not limiter.check(key_type):
if limit.result:
await MessageUtils.build_message(limit.result).send()
logger.debug(
f"{limit.module}({limit.limit_type}) 正在限制中...",
"AuthChecker",
session=session,
)
raise IgnoredException(f"{limit.module} 正在限制中...")
else:
logger.debug(
f"开始进行限制 {limit.module}({limit.limit_type})...",
"AuthChecker",
session=user_id,
group_id=group_id,
)
if isinstance(limiter, FreqLimiter):
limiter.start_cd(key_type)
if isinstance(limiter, UserBlockLimiter):
limiter.set_true(key_type)
if isinstance(limiter, CountLimiter):
limiter.increase(key_type)
class IsSuperuserException(Exception):
pass
class AuthChecker:
"""
权限检查
"""
def __init__(self):
check_notice_info_cd = Config.get_config("hook", "CHECK_NOTICE_INFO_CD")
if check_notice_info_cd is None or check_notice_info_cd < 0:
raise ValueError("模块: [hook], 配置项: [CHECK_NOTICE_INFO_CD] 为空或小于0")
self._flmt = FreqLimiter(check_notice_info_cd)
self._flmt_g = FreqLimiter(check_notice_info_cd)
self._flmt_s = FreqLimiter(check_notice_info_cd)
self._flmt_c = FreqLimiter(check_notice_info_cd)
def is_send_limit_message(self, plugin: PluginInfo, sid: str) -> bool:
"""是否发送提示消息
参数:
plugin: PluginInfo
返回:
bool: 是否发送提示消息
"""
if not base_config.get("IS_SEND_TIP_MESSAGE"):
return False
if plugin.plugin_type == PluginType.DEPENDANT:
return False
return plugin.module != "ai" if self._flmt_s.check(sid) else False
async def auth(
self,
matcher: Matcher,
event: Event,
bot: Bot,
session: EventSession,
message: UniMsg,
):
"""权限检查
参数:
matcher: matcher
bot: bot
session: EventSession
message: UniMsg
"""
is_ignore = False
cost_gold = 0
user_id = session.id1
group_id = session.id3
channel_id = session.id2
if not group_id:
group_id = channel_id
channel_id = None
if matcher.type == "notice" and not isinstance(event, PokeNotifyEvent):
"""过滤除poke外的notice"""
return
if user_id and matcher.plugin and (module_path := matcher.plugin.module_name):
try:
user = await UserConsole.get_user(user_id, session.platform)
except IntegrityError as e:
logger.debug(
"重复创建用户，已跳过该次权限...",
"AuthChecker",
session=session,
e=e,
)
return
if plugin := await Cache.get(CacheType.PLUGINS, module_path):
if plugin.plugin_type == PluginType.HIDDEN:
logger.debug(
f"插件: {plugin.name}:{plugin.module} "
"为HIDDEN，已跳过权限检查..."
)
return
try:
cost_gold = await self.auth_cost(user, plugin, session)
if session.id1 in bot.config.superusers:
if plugin.plugin_type == PluginType.SUPERUSER:
raise IsSuperuserException()
if not plugin.limit_superuser:
cost_gold = 0
raise IsSuperuserException()
await self.auth_bot(plugin, bot.self_id)
await self.auth_group(plugin, session, message)
await self.auth_admin(plugin, session)
await self.auth_plugin(plugin, session, event)
await self.auth_limit(plugin, session)
except IsSuperuserException:
logger.debug(
"超级用户或被ban跳过权限检测...", "AuthChecker", session=session
)
except IgnoredException:
is_ignore = True
LimitManage.unblock(
matcher.plugin.name, user_id, group_id, channel_id
)
except AssertionError as e:
is_ignore = True
logger.debug("消息无法发送", session=session, e=e)
if cost_gold and user_id:
"""花费金币"""
try:
await UserConsole.reduce_gold(
user_id,
cost_gold,
GoldHandle.PLUGIN,
matcher.plugin.name if matcher.plugin else "",
session.platform,
)
except InsufficientGold:
if u := await UserConsole.get_user(user_id):
u.gold = 0
await u.save(update_fields=["gold"])
logger.debug(
f"调用功能花费金币: {cost_gold}", "AuthChecker", session=session
)
if is_ignore:
raise IgnoredException("权限检测 ignore")
async def auth_bot(self, plugin: PluginInfo, bot_id: str):
"""机器人权限
参数:
plugin: PluginInfo
bot_id: bot_id
"""
if not await BotConsole.get_bot_status(bot_id):
logger.debug("Bot休眠中阻断权限检测...", "AuthChecker")
raise IgnoredException("BotConsole休眠权限检测 ignore")
if await BotConsole.is_block_plugin(bot_id, plugin.module):
logger.debug(
f"Bot插件 {plugin.name}({plugin.module}) 权限检查结果为关闭...",
"AuthChecker",
)
raise IgnoredException("BotConsole插件权限检测 ignore")
async def auth_limit(self, plugin: PluginInfo, session: EventSession):
"""插件限制
参数:
plugin: PluginInfo
session: EventSession
"""
user_id = session.id1
group_id = session.id3
channel_id = session.id2
if not group_id:
group_id = channel_id
channel_id = None
if plugin.module not in LimitManage.add_module:
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
)
async def auth_plugin(
self, plugin: PluginInfo, session: EventSession, event: Event
):
"""插件状态
参数:
plugin: PluginInfo
session: EventSession
"""
group_id = session.id3
channel_id = session.id2
if not group_id:
group_id = channel_id
channel_id = None
if user_id := session.id1:
is_poke = isinstance(event, PokeNotifyEvent)
if group_id:
sid = group_id or user_id
if await GroupConsole.is_superuser_block_plugin(
group_id, plugin.module
):
"""超级用户群组插件状态"""
if self.is_send_limit_message(plugin, sid) and not is_poke:
self._flmt_s.start_cd(group_id or user_id)
await MessageUtils.build_message(
"超级管理员禁用了该群此功能..."
).send(reply_to=True)
logger.debug(
f"{plugin.name}({plugin.module}) 超级管理员禁用了该群此功能...",
"AuthChecker",
session=session,
)
raise IgnoredException("超级管理员禁用了该群此功能...")
if await GroupConsole.is_normal_block_plugin(group_id, plugin.module):
"""群组插件状态"""
if self.is_send_limit_message(plugin, sid) and not is_poke:
self._flmt_s.start_cd(group_id or user_id)
await MessageUtils.build_message("该群未开启此功能...").send(
reply_to=True
)
logger.debug(
f"{plugin.name}({plugin.module}) 未开启此功能...",
"AuthChecker",
session=session,
)
raise IgnoredException("该群未开启此功能...")
if plugin.block_type == BlockType.GROUP:
"""全局群组禁用"""
try:
if self.is_send_limit_message(plugin, sid) and not is_poke:
self._flmt_c.start_cd(group_id)
await MessageUtils.build_message(
"该功能在群组中已被禁用..."
).send(reply_to=True)
except Exception as e:
logger.error(
"auth_plugin 发送消息失败",
"AuthChecker",
session=session,
e=e,
)
logger.debug(
f"{plugin.name}({plugin.module}) 该插件在群组中已被禁用...",
"AuthChecker",
session=session,
)
raise IgnoredException("该插件在群组中已被禁用...")
else:
sid = user_id
if plugin.block_type == BlockType.PRIVATE:
"""全局私聊禁用"""
try:
if self.is_send_limit_message(plugin, sid) and not is_poke:
self._flmt_c.start_cd(user_id)
await MessageUtils.build_message(
"该功能在私聊中已被禁用..."
).send()
except Exception as e:
logger.error(
"auth_admin 发送消息失败",
"AuthChecker",
session=session,
e=e,
)
logger.debug(
f"{plugin.name}({plugin.module}) 该插件在私聊中已被禁用...",
"AuthChecker",
session=session,
)
raise IgnoredException("该插件在私聊中已被禁用...")
if not plugin.status and plugin.block_type == BlockType.ALL:
"""全局状态"""
if group_id and await GroupConsole.is_super_group(group_id):
raise IsSuperuserException()
logger.debug(
f"{plugin.name}({plugin.module}) 全局未开启此功能...",
"AuthChecker",
session=session,
)
if self.is_send_limit_message(plugin, sid) and not is_poke:
self._flmt_s.start_cd(group_id or user_id)
await MessageUtils.build_message("全局未开启此功能...").send()
raise IgnoredException("全局未开启此功能...")
async def auth_admin(self, plugin: PluginInfo, session: EventSession):
"""管理员命令 个人权限
参数:
plugin: PluginInfo
session: EventSession
"""
user_id = session.id1
if user_id and plugin.admin_level:
if group_id := session.id3 or session.id2:
if not await LevelUser.check_level(
user_id, group_id, plugin.admin_level
):
try:
if self._flmt.check(user_id):
self._flmt.start_cd(user_id)
await MessageUtils.build_message(
[
At(flag="user", target=user_id),
f"你的权限不足喔，"
f"该功能需要的权限等级: {plugin.admin_level}",
]
).send(reply_to=True)
except Exception as e:
logger.error(
"auth_admin 发送消息失败",
"AuthChecker",
session=session,
e=e,
)
logger.debug(
f"{plugin.name}({plugin.module}) 管理员权限不足...",
"AuthChecker",
session=session,
)
raise IgnoredException("管理员权限不足...")
elif not await LevelUser.check_level(user_id, None, plugin.admin_level):
try:
await MessageUtils.build_message(
f"你的权限不足喔，该功能需要的权限等级: {plugin.admin_level}"
).send()
except Exception as e:
logger.error(
"auth_admin 发送消息失败", "AuthChecker", session=session, e=e
)
logger.debug(
f"{plugin.name}({plugin.module}) 管理员权限不足...",
"AuthChecker",
session=session,
)
raise IgnoredException("权限不足")
async def auth_group(
self, plugin: PluginInfo, session: EventSession, message: UniMsg
):
"""群黑名单检测 群总开关检测
参数:
plugin: PluginInfo
session: EventSession
message: UniMsg
"""
if not (group_id := session.id3 or session.id2):
return
text = message.extract_plain_text()
group = await GroupConsole.get_group(group_id)
if not group:
"""群不存在"""
logger.debug(
"群组信息不存在...",
"AuthChecker",
session=session,
)
raise IgnoredException("群不存在")
if group.level < 0:
"""群权限小于0"""
logger.debug(
"群黑名单, 群权限-1...",
"AuthChecker",
session=session,
)
raise IgnoredException("群黑名单")
if not group.status:
"""群休眠"""
if text.strip() != "醒来":
logger.debug("群休眠状态...", "AuthChecker", session=session)
raise IgnoredException("群休眠状态")
if plugin.level > group.level:
"""插件等级大于群等级"""
logger.debug(
f"{plugin.name}({plugin.module}) 群等级限制.."
f"该功能需要的群等级: {plugin.level}..",
"AuthChecker",
session=session,
)
raise IgnoredException(f"{plugin.name}({plugin.module}) 群等级限制...")
async def auth_cost(
self, user: UserConsole, plugin: PluginInfo, session: EventSession
) -> int:
"""检测是否满足金币条件
参数:
user: UserConsole
plugin: PluginInfo
session: EventSession
返回:
int: 需要消耗的金币
"""
if user.gold < plugin.cost_gold:
"""插件消耗金币不足"""
try:
await MessageUtils.build_message(
f"金币不足..该功能需要{plugin.cost_gold}金币.."
).send()
except Exception as e:
logger.error(
"auth_cost 发送消息失败", "AuthChecker", session=session, e=e
)
logger.debug(
f"{plugin.name}({plugin.module}) 金币限制.."
f"该功能需要{plugin.cost_gold}金币..",
"AuthChecker",
session=session,
)
raise IgnoredException(f"{plugin.name}({plugin.module}) 金币限制...")
return plugin.cost_gold
checker = AuthChecker()

76
zhenxun/builtin_plugins/hooks/auth/auth_admin.py Normal file
View File

@ -0,0 +1,76 @@
from nonebot.exception import IgnoredException
from nonebot_plugin_alconna import At
from nonebot_plugin_uninfo import Uninfo
from zhenxun.models.level_user import LevelUser
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.services.log import logger
from zhenxun.utils.cache_utils import Cache
from zhenxun.utils.enum import CacheType
from zhenxun.utils.message import MessageUtils
from .utils import freq
async def auth_admin(plugin: PluginInfo, session: Uninfo):
"""管理员命令 个人权限
参数:
plugin: PluginInfo
session: PluginInfo
"""
group_id = None
cache = Cache[list[LevelUser]](CacheType.LEVEL)
user_level = await cache.get(session.user.id) or []
if session.group:
if session.group.parent:
group_id = session.group.parent.id
else:
group_id = session.group.id
if not plugin.admin_level:
return
if group_id:
user_level += await cache.get(session.user.id, group_id) or []
user = max(user_level, key=lambda x: x.user_level)
if user.user_level < plugin.admin_level:
try:
if freq._flmt.check(session.user.id):
freq._flmt.start_cd(session.user.id)
await MessageUtils.build_message(
[
At(flag="user", target=session.user.id),
"你的权限不足喔，"
f"该功能需要的权限等级: {plugin.admin_level}",
]
).send(reply_to=True)
except Exception as e:
logger.error(
"auth_admin 发送消息失败",
"AuthChecker",
session=session,
e=e,
)
logger.debug(
f"{plugin.name}({plugin.module}) 管理员权限不足...",
"AuthChecker",
session=session,
)
raise IgnoredException("管理员权限不足...")
elif user_level:
user = max(user_level, key=lambda x: x.user_level)
if user.user_level < plugin.admin_level:
try:
await MessageUtils.build_message(
f"你的权限不足喔，该功能需要的权限等级: {plugin.admin_level}"
).send()
except Exception as e:
logger.error(
"auth_admin 发送消息失败", "AuthChecker", session=session, e=e
)
logger.debug(
f"{plugin.name}({plugin.module}) 管理员权限不足...",
"AuthChecker",
session=session,
)
raise IgnoredException("权限不足")

38
zhenxun/builtin_plugins/hooks/auth/auth_bot.py
View File

@ -4,31 +4,29 @@ from zhenxun.models.bot_console import BotConsole
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.services.log import logger
from zhenxun.utils.cache_utils import Cache
from zhenxun.utils.common_utils import CommonUtils
from zhenxun.utils.enum import CacheType
async def get_bot_status(bot_id: str):
a = await Cache.get(CacheType.BOT, bot_id)
cache_data = await Cache.get_cache(CacheType.BOT)
if cache_data and cache_data.getter:
b = await cache_data.getter.get(cache_data.data)
if bot := await Cache.get(CacheType.BOT, bot_id):
return bot
async def auth_bot(plugin: PluginInfo, bot_id: str):
"""机器人权限
"""bot层面的权限检查
参数:
plugin: PluginInfo
bot_id: bot_id
bot_id: bot id
异常:
IgnoredException: 忽略插件
IgnoredException: 忽略插件
"""
if not await BotConsole.get_bot_status(bot_id):
logger.debug("Bot休眠中阻断权限检测...", "AuthChecker")
raise IgnoredException("BotConsole休眠权限检测 ignore")
if await BotConsole.is_block_plugin(bot_id, plugin.module):
logger.debug(
f"Bot插件 {plugin.name}({plugin.module}) 权限检查结果为关闭...",
"AuthChecker",
)
raise IgnoredException("BotConsole插件权限检测 ignore")
if cache := Cache[BotConsole](CacheType.BOT):
bot = await cache.get(bot_id)
if not bot or not bot.status:
logger.debug("Bot不存在或休眠中阻断权限检测...", "AuthChecker")
raise IgnoredException("BotConsole休眠权限检测 ignore")
if CommonUtils.format(plugin.module) in bot.block_plugins:
logger.debug(
f"Bot插件 {plugin.name}({plugin.module}) 权限检查结果为关闭...",
"AuthChecker",
)
raise IgnoredException("BotConsole插件权限检测 ignore")

35
zhenxun/builtin_plugins/hooks/auth/auth_cost.py Normal file
View File

@ -0,0 +1,35 @@
from nonebot.exception import IgnoredException
from nonebot_plugin_uninfo import Uninfo
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.models.user_console import UserConsole
from zhenxun.services.log import logger
from zhenxun.utils.message import MessageUtils
async def auth_cost(user: UserConsole, plugin: PluginInfo, session: Uninfo) -> int:
"""检测是否满足金币条件
参数:
plugin: PluginInfo
session: Uninfo
返回:
int: 需要消耗的金币
"""
if user.gold < plugin.cost_gold:
"""插件消耗金币不足"""
try:
await MessageUtils.build_message(
f"金币不足..该功能需要{plugin.cost_gold}金币.."
).send()
except Exception as e:
logger.error("auth_cost 发送消息失败", "AuthChecker", session=session, e=e)
logger.debug(
f"{plugin.name}({plugin.module}) 金币限制.."
f"该功能需要{plugin.cost_gold}金币..",
"AuthChecker",
session=session,
)
raise IgnoredException(f"{plugin.name}({plugin.module}) 金币限制...")
return plugin.cost_gold

57
zhenxun/builtin_plugins/hooks/auth/auth_group.py Normal file
View File

@ -0,0 +1,57 @@
from nonebot.exception import IgnoredException
from nonebot_plugin_alconna import UniMsg
from nonebot_plugin_uninfo import Uninfo
from zhenxun.models.group_console import GroupConsole
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.services.log import logger
from zhenxun.utils.cache_utils import Cache
from zhenxun.utils.enum import CacheType
async def auth_group(plugin: PluginInfo, session: Uninfo, message: UniMsg):
"""群黑名单检测 群总开关检测
参数:
plugin: PluginInfo
session: EventSession
message: UniMsg
"""
if not session.group:
return
if session.group.parent:
group_id = session.group.parent.id
else:
group_id = session.group.id
text = message.extract_plain_text()
group = await Cache[GroupConsole](CacheType.GROUPS).get(group_id)
if not group:
"""群不存在"""
logger.debug(
"群组信息不存在...",
"AuthChecker",
session=session,
)
raise IgnoredException("群不存在")
if group.level < 0:
"""群权限小于0"""
logger.debug(
"群黑名单, 群权限-1...",
"AuthChecker",
session=session,
)
raise IgnoredException("群黑名单")
if not group.status:
"""群休眠"""
if text.strip() != "醒来":
logger.debug("群休眠状态...", "AuthChecker", session=session)
raise IgnoredException("群休眠状态")
if plugin.level > group.level:
"""插件等级大于群等级"""
logger.debug(
f"{plugin.name}({plugin.module}) 群等级限制.."
f"该功能需要的群等级: {plugin.level}..",
"AuthChecker",
session=session,
)
raise IgnoredException(f"{plugin.name}({plugin.module}) 群等级限制...")

189
zhenxun/builtin_plugins/hooks/auth/auth_limit.py Normal file
View File

@ -0,0 +1,189 @@
from typing import ClassVar
from nonebot.exception import IgnoredException
from nonebot_plugin_uninfo import Uninfo
from pydantic import BaseModel
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.models.plugin_limit import PluginLimit
from zhenxun.services.log import logger
from zhenxun.utils.enum import LimitWatchType, PluginLimitType
from zhenxun.utils.message import MessageUtils
from zhenxun.utils.utils import CountLimiter, FreqLimiter, UserBlockLimiter
class Limit(BaseModel):
limit: PluginLimit
limiter: FreqLimiter | UserBlockLimiter | CountLimiter
class Config:
arbitrary_types_allowed = True
class LimitManage:
add_module: ClassVar[list] = []
cd_limit: ClassVar[dict[str, Limit]] = {}
block_limit: ClassVar[dict[str, Limit]] = {}
count_limit: ClassVar[dict[str, Limit]] = {}
@classmethod
def add_limit(cls, limit: PluginLimit):
"""添加限制
参数:
limit: PluginLimit
"""
if limit.module not in cls.add_module:
cls.add_module.append(limit.module)
if limit.limit_type == PluginLimitType.BLOCK:
cls.block_limit[limit.module] = Limit(
limit=limit, limiter=UserBlockLimiter()
)
elif limit.limit_type == PluginLimitType.CD:
cls.cd_limit[limit.module] = Limit(
limit=limit, limiter=FreqLimiter(limit.cd)
)
elif limit.limit_type == PluginLimitType.COUNT:
cls.count_limit[limit.module] = Limit(
limit=limit, limiter=CountLimiter(limit.max_count)
)
@classmethod
def unblock(
cls, module: str, user_id: str, group_id: str | None, channel_id: str | None
):
"""解除插件block
参数:
module: 模块名
user_id: 用户id
group_id: 群组id
channel_id: 频道id
"""
if limit_model := cls.block_limit.get(module):
limit = limit_model.limit
limiter: UserBlockLimiter = limit_model.limiter # type: ignore
key_type = user_id
if group_id and limit.watch_type == LimitWatchType.GROUP:
key_type = channel_id or group_id
logger.debug(
f"解除对象: {key_type} 的block限制",
"AuthChecker",
session=user_id,
group_id=group_id,
)
limiter.set_false(key_type)
@classmethod
async def check(
cls,
module: str,
user_id: str,
group_id: str | None,
channel_id: str | None,
session: Uninfo,
):
"""检测限制
参数:
module: 模块名
user_id: 用户id
group_id: 群组id
channel_id: 频道id
session: Session
异常:
IgnoredException: IgnoredException
"""
if limit_model := cls.cd_limit.get(module):
await cls.__check(limit_model, user_id, group_id, channel_id, session)
if limit_model := cls.block_limit.get(module):
await cls.__check(limit_model, user_id, group_id, channel_id, session)
if limit_model := cls.count_limit.get(module):
await cls.__check(limit_model, user_id, group_id, channel_id, session)
@classmethod
async def __check(
cls,
limit_model: Limit | None,
user_id: str,
group_id: str | None,
channel_id: str | None,
session: Uninfo,
):
"""检测限制
参数:
limit_model: Limit
user_id: 用户id
group_id: 群组id
channel_id: 频道id
session: Session
异常:
IgnoredException: IgnoredException
"""
if not limit_model:
return
limit = limit_model.limit
limiter = limit_model.limiter
is_limit = (
LimitWatchType.ALL
or (group_id and limit.watch_type == LimitWatchType.GROUP)
or (not group_id and limit.watch_type == LimitWatchType.USER)
)
key_type = user_id
if group_id and limit.watch_type == LimitWatchType.GROUP:
key_type = channel_id or group_id
if is_limit and not limiter.check(key_type):
if limit.result:
await MessageUtils.build_message(limit.result).send()
logger.debug(
f"{limit.module}({limit.limit_type}) 正在限制中...",
"AuthChecker",
session=session,
)
raise IgnoredException(f"{limit.module} 正在限制中...")
else:
logger.debug(
f"开始进行限制 {limit.module}({limit.limit_type})...",
"AuthChecker",
session=user_id,
group_id=group_id,
)
if isinstance(limiter, FreqLimiter):
limiter.start_cd(key_type)
if isinstance(limiter, UserBlockLimiter):
limiter.set_true(key_type)
if isinstance(limiter, CountLimiter):
limiter.increase(key_type)
async def auth_limit(plugin: PluginInfo, session: Uninfo):
"""插件限制
参数:
plugin: PluginInfo
session: Uninfo
"""
user_id = session.user.id
group_id = None
channel_id = None
if session.group:
if session.group.parent:
group_id = session.group.parent.id
channel_id = session.group.id
else:
group_id = session.group.id
if not group_id:
group_id = channel_id
channel_id = None
if plugin.module not in LimitManage.add_module:
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)

201
zhenxun/builtin_plugins/hooks/auth/auth_plugin.py Normal file
View File

@ -0,0 +1,201 @@
from nonebot.adapters import Event
from nonebot.exception import IgnoredException
from nonebot_plugin_uninfo import Uninfo
from zhenxun.models.group_console import GroupConsole
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.services.log import logger
from zhenxun.utils.cache_utils import Cache
from zhenxun.utils.common_utils import CommonUtils
from zhenxun.utils.enum import BlockType, CacheType
from zhenxun.utils.message import MessageUtils
from .exception import IsSuperuserException
from .utils import freq
class GroupCheck:
def __init__(
self, plugin: PluginInfo, group_id: str, session: Uninfo, is_poke: bool
) -> None:
self.group_id = group_id
self.session = session
self.is_poke = is_poke
self.plugin = plugin
async def __get_data(self):
cache = Cache[GroupConsole](CacheType.GROUPS)
return await cache.get(self.group_id)
async def check(self):
await self.check_superuser_block(self.plugin)
async def check_superuser_block(self, plugin: PluginInfo):
"""超级用户禁用群组插件检测
参数:
plugin: PluginInfo
异常:
IgnoredException: 忽略插件
"""
group = await self.__get_data()
if group and CommonUtils.format(plugin.module) in group.superuser_block_plugin:
if freq.is_send_limit_message(plugin, group.group_id, self.is_poke):
freq._flmt_s.start_cd(self.group_id)
await MessageUtils.build_message("超级管理员禁用了该群此功能...").send(
reply_to=True
)
logger.debug(
f"{plugin.name}({plugin.module}) 超级管理员禁用了该群此功能...",
"AuthChecker",
session=self.session,
)
raise IgnoredException("超级管理员禁用了该群此功能...")
await self.check_normal_block(self.plugin)
async def check_normal_block(self, plugin: PluginInfo):
"""群组插件状态
参数:
plugin: PluginInfo
异常:
IgnoredException: 忽略插件
"""
group = await self.__get_data()
if group and CommonUtils.format(plugin.module) in group.block_plugin:
if freq.is_send_limit_message(plugin, self.group_id, self.is_poke):
freq._flmt_s.start_cd(self.group_id)
await MessageUtils.build_message("该群未开启此功能...").send(
reply_to=True
)
logger.debug(
f"{plugin.name}({plugin.module}) 未开启此功能...",
"AuthChecker",
session=self.session,
)
raise IgnoredException("该群未开启此功能...")
await self.check_global_block(self.plugin)
async def check_global_block(self, plugin: PluginInfo):
"""全局禁用插件检测
参数:
plugin: PluginInfo
异常:
IgnoredException: 忽略插件
"""
if plugin.block_type == BlockType.GROUP:
"""全局群组禁用"""
try:
if freq.is_send_limit_message(plugin, self.group_id, self.is_poke):
freq._flmt_c.start_cd(self.group_id)
await MessageUtils.build_message("该功能在群组中已被禁用...").send(
reply_to=True
)
except Exception as e:
logger.error(
"auth_plugin 发送消息失败",
"AuthChecker",
session=self.session,
e=e,
)
logger.debug(
f"{plugin.name}({plugin.module}) 该插件在群组中已被禁用...",
"AuthChecker",
session=self.session,
)
raise IgnoredException("该插件在群组中已被禁用...")
class PluginCheck:
def __init__(self, group_id: str | None, session: Uninfo, is_poke: bool):
self.session = session
self.is_poke = is_poke
self.group_id = group_id
async def check_user(self, plugin: PluginInfo):
"""全局私聊禁用检测
参数:
plugin: PluginInfo
异常:
IgnoredException: 忽略插件
"""
if plugin.block_type == BlockType.PRIVATE:
try:
if freq.is_send_limit_message(
plugin, self.session.user.id, self.is_poke
):
freq._flmt_c.start_cd(self.session.user.id)
await MessageUtils.build_message("该功能在私聊中已被禁用...").send()
except Exception as e:
logger.error(
"auth_admin 发送消息失败",
"AuthChecker",
session=self.session,
e=e,
)
logger.debug(
f"{plugin.name}({plugin.module}) 该插件在私聊中已被禁用...",
"AuthChecker",
session=self.session,
)
raise IgnoredException("该插件在私聊中已被禁用...")
async def check_global(self, plugin: PluginInfo):
"""全局状态
参数:
plugin: PluginInfo
异常:
IgnoredException: 忽略插件
"""
if not plugin.status and plugin.block_type == BlockType.ALL:
"""全局状态"""
cache = Cache[GroupConsole](CacheType.GROUPS)
if self.group_id and (group := await cache.get(self.group_id)):
if group.is_super:
raise IsSuperuserException()
logger.debug(
f"{plugin.name}({plugin.module}) 全局未开启此功能...",
"AuthChecker",
session=self.session,
)
sid = self.group_id or self.session.user.id
if freq.is_send_limit_message(plugin, sid, self.is_poke):
freq._flmt_s.start_cd(sid)
await MessageUtils.build_message("全局未开启此功能...").send()
raise IgnoredException("全局未开启此功能...")
async def auth_plugin(plugin: PluginInfo, session: Uninfo, event: Event):
"""插件状态
参数:
plugin: PluginInfo
session: Uninfo
"""
group_id = None
if session.group:
if session.group.parent:
group_id = session.group.parent.id
else:
group_id = session.group.id
try:
from nonebot.adapters.onebot.v11 import PokeNotifyEvent
is_poke = isinstance(event, PokeNotifyEvent)
except ImportError:
is_poke = False
user_check = PluginCheck(group_id, session, is_poke)
if group_id:
group_check = GroupCheck(plugin, group_id, session, is_poke)
await group_check.check()
else:
await user_check.check_user(plugin)
await user_check.check_global(plugin)

2
zhenxun/builtin_plugins/hooks/auth/exception.py Normal file
View File

@ -0,0 +1,2 @@
class IsSuperuserException(Exception):
pass

41
zhenxun/builtin_plugins/hooks/auth/utils.py Normal file
View File

@ -0,0 +1,41 @@
from zhenxun.configs.config import Config
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.utils.enum import PluginType
from zhenxun.utils.utils import FreqLimiter
base_config = Config.get("hook")
class FreqUtils:
def __init__(self):
check_notice_info_cd = Config.get_config("hook", "CHECK_NOTICE_INFO_CD")
if check_notice_info_cd is None or check_notice_info_cd < 0:
raise ValueError("模块: [hook], 配置项: [CHECK_NOTICE_INFO_CD] 为空或小于0")
self._flmt = FreqLimiter(check_notice_info_cd)
self._flmt_g = FreqLimiter(check_notice_info_cd)
self._flmt_s = FreqLimiter(check_notice_info_cd)
self._flmt_c = FreqLimiter(check_notice_info_cd)
def is_send_limit_message(
self, plugin: PluginInfo, sid: str, is_poke: bool
) -> bool:
"""是否发送提示消息
参数:
plugin: PluginInfo
sid: 检测键
is_poke: 是否是戳一戳
返回:
bool: 是否发送提示消息
"""
if is_poke:
return False
if not base_config.get("IS_SEND_TIP_MESSAGE"):
return False
if plugin.plugin_type == PluginType.DEPENDANT:
return False
return plugin.module != "ai" if self._flmt_s.check(sid) else False
freq = FreqUtils()

125
zhenxun/builtin_plugins/hooks/auth_checker.py Normal file
View File

@ -0,0 +1,125 @@
from nonebot.adapters import Bot, Event
from nonebot.exception import IgnoredException
from nonebot.matcher import Matcher
from nonebot_plugin_alconna import UniMsg
from nonebot_plugin_uninfo import Uninfo
from tortoise.exceptions import IntegrityError
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.models.user_console import UserConsole
from zhenxun.services.log import logger
from zhenxun.utils.cache_utils import Cache
from zhenxun.utils.enum import (
CacheType,
GoldHandle,
PluginType,
)
from zhenxun.utils.exception import InsufficientGold
from zhenxun.utils.platform import PlatformUtils
from .auth.auth_admin import auth_admin
from .auth.auth_bot import auth_bot
from .auth.auth_cost import auth_cost
from .auth.auth_group import auth_group
from .auth.auth_limit import LimitManage, auth_limit
from .auth.auth_plugin import auth_plugin
from .auth.exception import IsSuperuserException
async def auth(
matcher: Matcher,
event: Event,
bot: Bot,
session: Uninfo,
message: UniMsg,
):
"""权限检查
参数:
matcher: matcher
event: Event
bot: bot
session: Uninfo
message: UniMsg
"""
user_id = session.user.id
group_id = None
channel_id = None
if session.group:
if session.group.parent:
group_id = session.group.parent.id
channel_id = session.group.id
else:
group_id = session.group.id
is_ignore = False
cost_gold = 0
try:
from nonebot.adapters.onebot.v11 import PokeNotifyEvent
if matcher.type == "notice" and not isinstance(event, PokeNotifyEvent):
"""过滤除poke外的notice"""
return
except ImportError:
pass
user_cache = Cache[UserConsole](CacheType.USERS)
if matcher.plugin and (module := matcher.plugin.name):
try:
user = await user_cache.get(session.user.id)
except IntegrityError as e:
logger.debug(
"重复创建用户，已跳过该次权限检查...",
"AuthChecker",
session=session,
e=e,
)
return
plugin = await Cache[PluginInfo](CacheType.PLUGINS).get(module)
if user and plugin:
if plugin.plugin_type == PluginType.HIDDEN:
logger.debug(
f"插件: {plugin.name}:{plugin.module} "
"为HIDDEN，已跳过权限检查..."
)
return
try:
cost_gold = await auth_cost(user, plugin, session)
if session.user.id in bot.config.superusers:
if plugin.plugin_type == PluginType.SUPERUSER:
raise IsSuperuserException()
if not plugin.limit_superuser:
cost_gold = 0
raise IsSuperuserException()
await auth_bot(plugin, bot.self_id)
await auth_group(plugin, session, message)
await auth_admin(plugin, session)
await auth_plugin(plugin, session, event)
await auth_limit(plugin, session)
except IsSuperuserException:
logger.debug(
"超级用户或被ban跳过权限检测...", "AuthChecker", session=session
)
except IgnoredException:
is_ignore = True
LimitManage.unblock(matcher.plugin.name, user_id, group_id, channel_id)
except AssertionError as e:
is_ignore = True
logger.debug("消息无法发送", session=session, e=e)
if cost_gold and user_id:
"""花费金币"""
try:
await UserConsole.reduce_gold(
user_id,
cost_gold,
GoldHandle.PLUGIN,
matcher.plugin.name if matcher.plugin else "",
PlatformUtils.get_platform(session),
)
except InsufficientGold:
if u := await UserConsole.get_user(user_id):
u.gold = 0
await u.save(update_fields=["gold"])
# 更新缓存
await user_cache.update(user_id)
logger.debug(f"调用功能花费金币: {cost_gold}", "AuthChecker", session=session)
if is_ignore:
raise IgnoredException("权限检测 ignore")

35
zhenxun/builtin_plugins/hooks/auth_hook.py
View File

@ -1,18 +1,16 @@
from nonebot.adapters.onebot.v11 import Bot, Event
from nonebot.adapters import Bot, Event
from nonebot.matcher import Matcher
from nonebot.message import run_postprocessor, run_preprocessor
from nonebot_plugin_alconna import UniMsg
from nonebot_plugin_session import EventSession
from nonebot_plugin_uninfo import Uninfo
from ._auth_checker import LimitManage, checker
from .auth_checker import LimitManage, auth
# # 权限检测
@run_preprocessor
async def _(
matcher: Matcher, event: Event, bot: Bot, session: EventSession, message: UniMsg
):
await checker.auth(
async def _(matcher: Matcher, event: Event, bot: Bot, session: Uninfo, message: UniMsg):
await auth(
matcher,
event,
bot,
@ -23,19 +21,16 @@ async def _(
# 解除命令block阻塞
@run_postprocessor
async def _(
matcher: Matcher,
exception: Exception | None,
bot: Bot,
event: Event,
session: EventSession,
):
user_id = session.id1
group_id = session.id3
channel_id = session.id2
if not group_id:
group_id = channel_id
channel_id = None
async def _(matcher: Matcher, session: Uninfo):
user_id = session.user.id
group_id = None
channel_id = None
if session.group:
if session.group.parent:
group_id = session.group.parent.id
channel_id = session.group.id
else:
group_id = session.group.id
if user_id and matcher.plugin:
module = matcher.plugin.name
LimitManage.unblock(module, user_id, group_id, channel_id)

127
zhenxun/builtin_plugins/init/__init__.py
View File

@ -1,4 +1,5 @@
from pathlib import Path
from typing import Any
import nonebot
from nonebot.adapters import Bot
@ -6,9 +7,11 @@ from nonebot.adapters import Bot
from zhenxun.models.ban_console import BanConsole
from zhenxun.models.bot_console import BotConsole
from zhenxun.models.group_console import GroupConsole
from zhenxun.models.level_user import LevelUser
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.models.user_console import UserConsole
from zhenxun.services.log import logger
from zhenxun.utils.cache_utils import Cache
from zhenxun.utils.cache_utils import CacheRoot
from zhenxun.utils.enum import CacheType
from zhenxun.utils.platform import PlatformUtils
@ -48,45 +51,131 @@ async def _(bot: Bot):
)
@Cache.listener(CacheType.PLUGINS)
async def _():
data_list = await PluginInfo.get_plugins()
return {p.module: p for p in data_list}
@CacheRoot.new(CacheType.PLUGINS)
async def _(data: dict[str, PluginInfo] = {}, key: str | None = None):
if data and key:
if plugin := await PluginInfo.get_plugin(module=key):
data[key] = plugin
else:
data_list = await PluginInfo.get_plugins()
return {p.module: p for p in data_list}
@Cache.getter(CacheType.PLUGINS, result_model=PluginInfo)
def _(data: dict[str, PluginInfo], module: str):
return data.get(module, None)
@CacheRoot.updater(CacheType.PLUGINS)
async def _(data: dict[str, PluginInfo], key: str, value: Any):
if value:
data[key] = value
elif plugin := await PluginInfo.get_plugin(module=key):
data[key] = plugin
@Cache.listener(CacheType.GROUPS)
@CacheRoot.getter(CacheType.PLUGINS, result_model=PluginInfo)
async def _(data: dict[str, PluginInfo], module: str):
result = data.get(module, None)
if not result:
result = await PluginInfo.get_plugin(module=module)
if result:
data[module] = result
return result
@CacheRoot.new(CacheType.GROUPS)
async def _():
data_list = await GroupConsole.all()
return {p.group_id: p for p in data_list}
return {p.group_id: p for p in data_list if not p.channel_id}
@Cache.getter(CacheType.GROUPS, result_model=GroupConsole)
def _(data: dict[str, GroupConsole], module: str):
return data.get(module, None)
@CacheRoot.updater(CacheType.GROUPS)
async def _(data: dict[str, GroupConsole], key: str, value: Any):
if value:
data[key] = value
elif group := await GroupConsole.get_group(group_id=key):
data[key] = group
@Cache.listener(CacheType.BOT)
@CacheRoot.getter(CacheType.GROUPS, result_model=GroupConsole)
async def _(data: dict[str, GroupConsole], group_id: str):
result = data.get(group_id, None)
if not result:
result = await GroupConsole.get_group(group_id=group_id)
if result:
data[group_id] = result
return result
@CacheRoot.new(CacheType.BOT)
async def _():
data_list = await BotConsole.all()
return {p.bot_id: p for p in data_list}
@Cache.getter(CacheType.BOT, result_model=BotConsole)
def _(data: dict[str, BotConsole], module: str):
return data.get(module, None)
@CacheRoot.updater(CacheType.BOT)
async def _(data: dict[str, BotConsole], key: str, value: Any):
if value:
data[key] = value
elif bot := await BotConsole.get_or_none(bot_id=key):
data[key] = bot
@Cache.listener(CacheType.BAN)
@CacheRoot.getter(CacheType.BOT, result_model=BotConsole)
async def _(data: dict[str, BotConsole], bot_id: str):
result = data.get(bot_id, None)
if not result:
result = await BotConsole.get_or_none(bot_id=bot_id)
if result:
data[bot_id] = result
return result
@CacheRoot.new(CacheType.USERS)
async def _():
data_list = await UserConsole.all()
return {p.user_id: p for p in data_list}
@CacheRoot.updater(CacheType.USERS)
async def _(data: dict[str, UserConsole], key: str, value: Any):
if value:
data[key] = value
elif user := await UserConsole.get_user(user_id=key):
data[key] = user
@CacheRoot.getter(CacheType.USERS, result_model=UserConsole)
async def _(data: dict[str, UserConsole], user_id: str):
result = data.get(user_id, None)
if not result:
result = await UserConsole.get_user(user_id=user_id)
if result:
data[user_id] = result
return result
@CacheRoot.new(CacheType.LEVEL)
async def _():
return await LevelUser().all()
@CacheRoot.getter(CacheType.LEVEL, result_model=list[LevelUser])
def _(data_list: list[LevelUser], user_id: str, group_id: str | None = None):
if not group_id:
return [
data for data in data_list if data.user_id == user_id and not data.group_id
]
else:
return [
data
for data in data_list
if data.user_id == user_id and data.group_id == group_id
]
@CacheRoot.new(CacheType.BAN)
async def _():
return await BanConsole.all()
@Cache.getter(CacheType.BAN, result_model=list[BanConsole])
@CacheRoot.getter(CacheType.BAN, result_model=list[BanConsole])
def _(data_list: list[BanConsole], user_id: str, group_id: str):
if user_id:
if group_id:

13
zhenxun/models/group_console.py
View File

@ -7,7 +7,8 @@ from tortoise.backends.base.client import BaseDBAsyncClient
from zhenxun.models.plugin_info import PluginInfo
from zhenxun.models.task_info import TaskInfo
from zhenxun.services.db_context import Model
from zhenxun.utils.enum import PluginType
from zhenxun.utils.cache_utils import CacheRoot
from zhenxun.utils.enum import CacheType, PluginType
class GroupConsole(Model):
@ -46,8 +47,7 @@ class GroupConsole(Model):
platform = fields.CharField(255, default="qq", description="所属平台")
"""所属平台"""
class Meta: # type: ignore
table = "group_console"
class Meta: # type: ignore table = "group_console"
table_description = "群组信息表"
unique_together = ("group_id", "channel_id")
@ -80,6 +80,7 @@ class GroupConsole(Model):
return "".join(cls.format(item) for item in data)
@classmethod
@CacheRoot.listener(CacheType.GROUPS)
async def create(
cls, using_db: BaseDBAsyncClient | None = None, **kwargs: Any
) -> Self:
@ -100,6 +101,7 @@ class GroupConsole(Model):
return group
@classmethod
@CacheRoot.listener(CacheType.GROUPS)
async def get_or_create(
cls,
defaults: dict | None = None,
@ -127,6 +129,7 @@ class GroupConsole(Model):
return group, is_create
@classmethod
@CacheRoot.listener(CacheType.GROUPS)
async def update_or_create(
cls,
defaults: dict | None = None,
@ -216,6 +219,7 @@ class GroupConsole(Model):
)
@classmethod
@CacheRoot.listener(CacheType.GROUPS)
async def set_block_plugin(
cls,
group_id: str,
@ -242,6 +246,7 @@ class GroupConsole(Model):
await group.save(update_fields=["block_plugin", "superuser_block_plugin"])
@classmethod
@CacheRoot.listener(CacheType.GROUPS)
async def set_unblock_plugin(
cls,
group_id: str,
@ -338,6 +343,7 @@ class GroupConsole(Model):
)
@classmethod
@CacheRoot.listener(CacheType.GROUPS)
async def set_block_task(
cls,
group_id: str,
@ -364,6 +370,7 @@ class GroupConsole(Model):
await group.save(update_fields=["block_task", "superuser_block_task"])
@classmethod
@CacheRoot.listener(CacheType.GROUPS)
async def set_unblock_task(
cls,
group_id: str,

12
zhenxun/models/plugin_info.py
View File

@ -1,10 +1,13 @@
from typing import Any
from typing_extensions import Self
from tortoise import fields
from tortoise.backends.base.client import BaseDBAsyncClient
from zhenxun.models.plugin_limit import PluginLimit # noqa: F401
from zhenxun.services.db_context import Model
from zhenxun.utils.enum import BlockType, PluginType
from zhenxun.utils.cache_utils import CacheRoot
from zhenxun.utils.enum import BlockType, CacheType, PluginType
class PluginInfo(Model):
@ -79,6 +82,13 @@ class PluginInfo(Model):
"""
return await cls.filter(load_status=load_status, **kwargs).all()
@classmethod
@CacheRoot.listener(CacheType.PLUGINS)
async def create(
cls, using_db: BaseDBAsyncClient | None = None, **kwargs: Any
) -> Self:
return await super().create(using_db=using_db, **kwargs)
@classmethod
async def _run_script(cls):
return [

121
zhenxun/utils/cache_utils.py
View File

@ -1,11 +1,14 @@
from collections.abc import Callable
from functools import wraps
import time
from typing import Any, ClassVar, Generic, TypeVar, cast
from nonebot.utils import is_coroutine_callable
from pydantic import BaseModel
__all__ = ["Cache", "CacheData"]
from zhenxun.services.log import logger
__all__ = ["Cache", "CacheData", "CacheRoot"]
T = TypeVar("T")
@ -27,10 +30,14 @@ class CacheGetter(BaseModel, Generic[T]):
class CacheData(BaseModel):
name: str
"""缓存名称"""
func: Callable[..., Any]
"""更新方法"""
getter: CacheGetter | None = None
"""获取方法"""
updater: Callable[..., Any] | None = None
"""更新单个方法"""
data: Any = None
"""缓存数据"""
expire: int
@ -40,13 +47,36 @@ class CacheData(BaseModel):
reload_count: int = 0
"""更新次数"""
async def reload(self):
async def get(self, *args, **kwargs) -> Any:
"""获取单个缓存"""
if not self.getter:
return self.data
return await self.getter.get(self.data, *args, **kwargs)
async def update(self, key: str, value: Any = None, *args, **kwargs):
"""更新单个缓存"""
if not self.updater:
return logger.warning(
f"缓存类型 {self.name} 没有更新方法，无法更新", "CacheRoot"
)
if self.data:
if is_coroutine_callable(self.updater):
await self.updater(self.data, key, value, *args, **kwargs)
else:
self.updater(self.data, key, value, *args, **kwargs)
else:
logger.warning(f"缓存类型 {self.name} 为空，无法更新", "CacheRoot")
async def reload(self, *args, **kwargs):
"""更新缓存"""
self.data = (
await self.func() if is_coroutine_callable(self.func) else self.func()
await self.func(*args, **kwargs)
if is_coroutine_callable(self.func)
else self.func(*args, **kwargs)
)
self.reload_time = time.time()
self.reload_count += 1
logger.debug(f"缓存类型 {self.name} 更新全局缓存", "CacheRoot")
async def check_expire(self):
if time.time() - self.reload_time > self.expire or not self.reload_count:
@ -54,14 +84,54 @@ class CacheData(BaseModel):
class CacheManage:
"""全局缓存管理，减少数据库与网络请求查询次数
异常:
ValueError: 数据名称重复
ValueError: 数据不存在
"""
_data: ClassVar[dict[str, CacheData]] = {}
def listener(self, name: str, expire: int = 60 * 10):
def new(self, name: str, expire: int = 60 * 10):
def wrapper(func: Callable):
_name = name.upper()
if _name in self._data:
raise ValueError(f"DbCache 缓存数据 {name} 已存在...")
self._data[_name] = CacheData(func=func, expire=expire)
self._data[_name] = CacheData(name=_name, func=func, expire=expire)
return wrapper
def listener(self, name: str):
def decorator(func):
@wraps(func)
async def wrapper(*args, **kwargs):
try:
if is_coroutine_callable(func):
result = await func(*args, **kwargs)
else:
result = func(*args, **kwargs)
return result
finally:
cache_data = self._data.get(name.upper())
if cache_data:
await cache_data.reload()
logger.debug(
f"缓存类型 {name.upper()} 进行监听更新...", "CacheRoot"
)
return wrapper
return decorator
def updater(self, name: str):
def wrapper(func: Callable):
_name = name.upper()
if _name not in self._data:
raise ValueError(f"DbCache 缓存数据 {name} 不存在...")
self._data[_name].updater = func
return wrapper
@ -83,12 +153,12 @@ class CacheManage:
):
await self._data[name].reload()
async def get_cache_data(self, name: str) -> CacheData | None:
async def get_cache_data(self, name: str):
if cache := await self.get_cache(name):
return cache
return cache.data
return None
async def get_cache(self, name: str):
async def get_cache(self, name: str) -> CacheData | None:
name = name.upper()
cache = self._data.get(name)
if cache:
@ -96,18 +166,35 @@ class CacheManage:
return cache
return None
async def get(self, name: str, *args, **kwargs) -> T | None:
cache = self._data.get(name.upper())
async def get(self, name: str, *args, **kwargs):
cache = await self.get_cache(name.upper())
if cache:
return (
await cache.getter.get(*args, **kwargs) if cache.getter else cache.data
)
return await cache.get(*args, **kwargs) if cache.getter else cache.data
return None
async def reload(self, name: str):
cache = self._data.get(name.upper())
async def reload(self, name: str, *args, **kwargs):
cache = await self.get_cache(name.upper())
if cache:
await cache.reload()
await cache.reload(*args, **kwargs)
async def update(self, name: str, key: str, value: Any, *args, **kwargs):
cache = await self.get_cache(name.upper())
if cache:
await cache.update(key, value, *args, **kwargs)
Cache = CacheManage()
CacheRoot = CacheManage()
class Cache(Generic[T]):
def __init__(self, module: str):
self.module = module
async def get(self, *args, **kwargs) -> T | None:
return await CacheRoot.get(self.module, *args, **kwargs)
async def update(self, key: str, value: Any = None, *args, **kwargs):
return await CacheRoot.update(self.module, key, value, *args, **kwargs)
async def reload(self, key: str | None = None, *args, **kwargs):
await CacheRoot.reload(self.module, key, *args, **kwargs)

4
zhenxun/utils/enum.py
View File

@ -10,10 +10,14 @@ class CacheType(StrEnum):
"""全局全部插件"""
GROUPS = "GLOBAL_ALL_GROUPS"
"""全局全部群组"""
USERS = "GLOBAL_ALL_USERS"
"""全部用户"""
BAN = "GLOBAL_ALL_BAN"
"""全局ban列表"""
BOT = "GLOBAL_BOT"
"""全局bot信息"""
LEVEL = "GLOBAL_USER_LEVEL"
"""用户权限"""
class GoldHandle(StrEnum):