mirror of
https://github.com/zhenxun-org/zhenxun_bot.git
synced 2026-09-28 16:20:56 +08:00
263 lines
7.1 KiB
Python
263 lines
7.1 KiB
Python
import asyncio
|
||
import time
|
||
|
||
from nonebot.matcher import Matcher
|
||
from nonebot_plugin_alconna import At
|
||
from nonebot_plugin_uninfo import Uninfo
|
||
|
||
from zhenxun.configs.config import Config
|
||
from zhenxun.models.ban_console import BanConsole
|
||
from zhenxun.models.plugin_info import PluginInfo
|
||
from zhenxun.services.cache.cache_containers import CacheDict
|
||
from zhenxun.services.cache.runtime_cache import BanMemoryCache
|
||
from zhenxun.services.db_context import DB_TIMEOUT_SECONDS
|
||
from zhenxun.services.log import logger
|
||
from zhenxun.utils.enum import PluginType
|
||
from zhenxun.utils.utils import EntityIDs, get_entity_ids
|
||
|
||
from .config import LOGGER_COMMAND, WARNING_THRESHOLD
|
||
from .exception import SkipPluginException
|
||
from .utils import freq
|
||
|
||
Config.add_plugin_config(
|
||
"hook",
|
||
"BAN_RESULT",
|
||
"才不会给你发消息.",
|
||
help="对被ban用户发送的消息",
|
||
)
|
||
BAN_CACHE_TTL = 2
|
||
BAN_CACHE_TTL_POSITIVE = 30
|
||
BAN_CACHE_TTL_NEGATIVE = 5
|
||
|
||
BAN_CACHE = (
|
||
CacheDict("AUTH_BAN_CACHE", expire=0)
|
||
if max(BAN_CACHE_TTL_POSITIVE, BAN_CACHE_TTL_NEGATIVE) > 0
|
||
else None
|
||
)
|
||
|
||
|
||
def _ban_cache_key(user_id: str | None, group_id: str | None) -> str:
|
||
return f"{user_id or ''}:{group_id or ''}"
|
||
|
||
|
||
def _ban_cache_get(key: str) -> int | None:
|
||
if not BAN_CACHE:
|
||
return None
|
||
try:
|
||
return BAN_CACHE[key]
|
||
except KeyError:
|
||
return None
|
||
|
||
|
||
def _ban_cache_set(key: str, value: int) -> None:
|
||
if not BAN_CACHE:
|
||
return
|
||
ttl = BAN_CACHE_TTL_POSITIVE if value else BAN_CACHE_TTL_NEGATIVE
|
||
if ttl <= 0:
|
||
return
|
||
BAN_CACHE.set(key, value, expire=ttl)
|
||
|
||
|
||
async def calculate_ban_time(ban_record: BanConsole | None) -> int:
|
||
"""根据ban记录计算剩余ban时间
|
||
|
||
参数:
|
||
ban_record: BanConsole记录
|
||
|
||
返回:
|
||
int: ban剩余时长，-1时为永久ban，0表示未被ban
|
||
"""
|
||
if not ban_record:
|
||
return 0
|
||
|
||
if ban_record.duration == -1:
|
||
return -1
|
||
|
||
_time = time.time() - (ban_record.ban_time + ban_record.duration)
|
||
if _time < 0:
|
||
return int(abs(_time))
|
||
await ban_record.delete()
|
||
return 0
|
||
|
||
|
||
async def is_ban(user_id: str | None, group_id: str | None) -> int:
|
||
"""检查用户或群组是否被ban
|
||
|
||
参数:
|
||
user_id: 用户ID
|
||
group_id: 群组ID
|
||
|
||
返回:
|
||
int: ban剩余时长，-1时为永久ban，0表示未被ban
|
||
"""
|
||
if not user_id and not group_id:
|
||
return 0
|
||
if not BanMemoryCache.is_loaded():
|
||
return 0
|
||
return BanMemoryCache.remaining_time(user_id, group_id)
|
||
|
||
|
||
def check_plugin_type(matcher: Matcher) -> bool:
|
||
"""判断插件类型是否是隐藏插件
|
||
|
||
参数:
|
||
matcher: Matcher
|
||
|
||
返回:
|
||
bool: 是否为隐藏插件
|
||
"""
|
||
if plugin := matcher.plugin:
|
||
if metadata := plugin.metadata:
|
||
extra = metadata.extra
|
||
if extra.get("plugin_type") in [PluginType.HIDDEN]:
|
||
return False
|
||
return True
|
||
|
||
|
||
def format_time(time_val: float) -> str:
|
||
"""格式化时间
|
||
|
||
参数:
|
||
time_val: ban时长
|
||
|
||
返回:
|
||
str: 格式化时间文本
|
||
"""
|
||
if time_val == -1:
|
||
return "∞"
|
||
time_val = abs(int(time_val))
|
||
if time_val < 60:
|
||
time_str = f"{time_val!s} 秒"
|
||
else:
|
||
minute = int(time_val / 60)
|
||
if minute > 60:
|
||
hours = minute // 60
|
||
minute %= 60
|
||
time_str = f"{hours} 小时 {minute}分钟"
|
||
else:
|
||
time_str = f"{minute} 分钟"
|
||
return time_str
|
||
|
||
|
||
async def group_handle(group_id: str) -> None:
|
||
"""群组ban检查
|
||
|
||
参数:
|
||
group_id: 群组id
|
||
|
||
异常:
|
||
SkipPluginException: 群组处于黑名单
|
||
"""
|
||
start_time = time.time()
|
||
try:
|
||
if await is_ban(None, group_id):
|
||
raise SkipPluginException("群组处于黑名单中...")
|
||
finally:
|
||
# 记录执行时间
|
||
elapsed = time.time() - start_time
|
||
if elapsed > WARNING_THRESHOLD: # 记录耗时超过500ms的检查
|
||
logger.warning(
|
||
f"group_handle 耗时: {elapsed:.3f}s",
|
||
LOGGER_COMMAND,
|
||
group_id=group_id,
|
||
)
|
||
|
||
|
||
async def user_handle(plugin: PluginInfo, entity: EntityIDs, session: Uninfo) -> None:
|
||
"""用户ban检查
|
||
|
||
参数:
|
||
module: 插件模块名
|
||
entity: 实体ID信息
|
||
session: Uninfo
|
||
|
||
异常:
|
||
SkipPluginException: 用户处于黑名单
|
||
"""
|
||
start_time = time.time()
|
||
try:
|
||
ban_result = Config.get_config("hook", "BAN_RESULT")
|
||
time_val = await is_ban(entity.user_id, entity.group_id)
|
||
if not time_val:
|
||
return
|
||
time_str = format_time(time_val)
|
||
|
||
if (
|
||
plugin
|
||
and time_val != -1
|
||
and ban_result
|
||
and freq.is_send_limit_message(plugin, entity.user_id, False)
|
||
):
|
||
raise SkipPluginException(
|
||
"用户处于黑名单中...",
|
||
tip_message=[
|
||
At(flag="user", target=entity.user_id),
|
||
f"{ban_result}\n在..在 {time_str} 后才会理你喔",
|
||
],
|
||
tip_check_tag=entity.user_id,
|
||
tip_timeout=DB_TIMEOUT_SECONDS,
|
||
)
|
||
raise SkipPluginException("用户处于黑名单中...")
|
||
finally:
|
||
# 记录执行时间
|
||
elapsed = time.time() - start_time
|
||
if elapsed > WARNING_THRESHOLD: # 记录耗时超过500ms的检查
|
||
logger.warning(
|
||
f"user_handle 耗时: {elapsed:.3f}s",
|
||
LOGGER_COMMAND,
|
||
session=session,
|
||
)
|
||
|
||
|
||
async def auth_ban(
|
||
matcher: Matcher,
|
||
session: Uninfo,
|
||
plugin: PluginInfo,
|
||
*,
|
||
entity: EntityIDs | None = None,
|
||
is_superuser: bool = False,
|
||
) -> None:
|
||
"""权限检查 - ban 检查
|
||
|
||
参数:
|
||
matcher: Matcher
|
||
session: Uninfo
|
||
"""
|
||
start_time = time.time()
|
||
try:
|
||
if not check_plugin_type(matcher):
|
||
return
|
||
if not matcher.plugin_name:
|
||
return
|
||
if entity is None:
|
||
entity = get_entity_ids(session)
|
||
if is_superuser:
|
||
return
|
||
if entity.group_id:
|
||
try:
|
||
await asyncio.wait_for(
|
||
group_handle(entity.group_id), timeout=DB_TIMEOUT_SECONDS
|
||
)
|
||
except asyncio.TimeoutError:
|
||
logger.error(f"群组ban检查超时: {entity.group_id}", LOGGER_COMMAND)
|
||
# 超时时不阻塞，继续执行
|
||
|
||
if entity.user_id:
|
||
try:
|
||
await asyncio.wait_for(
|
||
user_handle(plugin, entity, session),
|
||
timeout=DB_TIMEOUT_SECONDS,
|
||
)
|
||
except asyncio.TimeoutError:
|
||
logger.error(f"用户ban检查超时: {entity.user_id}", LOGGER_COMMAND)
|
||
# 超时时不阻塞，继续执行
|
||
finally:
|
||
# 记录总执行时间
|
||
elapsed = time.time() - start_time
|
||
if elapsed > WARNING_THRESHOLD: # 记录耗时超过500ms的检查
|
||
logger.warning(
|
||
f"auth_ban 总耗时: {elapsed:.3f}s, plugin={matcher.plugin_name}",
|
||
LOGGER_COMMAND,
|
||
session=session,
|
||
)
|