Dev (#26)
* feat: 整合开发历史 * codepy安全性升级 * 优化一些东西 * 再次优化 * 更新一下 requirements.txt * CQ码支持以及视频解析 * hotfix * 更新DEV readme.md * feat: 添加Docker沙箱代码执行功能 - 新增Docker沙箱执行环境,提供安全隔离的代码执行能力 - 重构code_py插件,使用Docker容器替代子进程执行 - 添加docker配置项和权限检查功能 - 实现代码执行队列和并发控制 - 新增广播插件,仅限管理员使用
This commit is contained in:
@@ -69,7 +69,7 @@ class CommandManager:
|
||||
|
||||
def command(
|
||||
self,
|
||||
name: str,
|
||||
*names: str,
|
||||
permission: Optional[Any] = None,
|
||||
override_permission_check: bool = False
|
||||
) -> Callable:
|
||||
@@ -77,7 +77,7 @@ class CommandManager:
|
||||
装饰器:注册一个消息指令处理器。
|
||||
"""
|
||||
return self.message_handler.command(
|
||||
name,
|
||||
*names,
|
||||
permission=permission,
|
||||
override_permission_check=override_permission_check
|
||||
)
|
||||
|
||||
@@ -73,6 +73,15 @@ class Config:
|
||||
"""
|
||||
return self._data.get("redis", {})
|
||||
|
||||
@property
|
||||
def docker(self) -> dict:
|
||||
"""
|
||||
获取 Docker 配置
|
||||
|
||||
:return: 配置字典
|
||||
"""
|
||||
return self._data.get("docker", {})
|
||||
|
||||
|
||||
# 实例化全局配置对象
|
||||
global_config = Config()
|
||||
|
||||
@@ -83,7 +83,7 @@ class MessageHandler(BaseHandler):
|
||||
|
||||
def command(
|
||||
self,
|
||||
name: str,
|
||||
*names: str,
|
||||
permission: Optional[Permission] = None,
|
||||
override_permission_check: bool = False
|
||||
) -> Callable:
|
||||
@@ -93,11 +93,12 @@ class MessageHandler(BaseHandler):
|
||||
def decorator(func: Callable) -> Callable:
|
||||
if not inspect.iscoroutinefunction(func):
|
||||
raise SyncHandlerError(f"命令处理器 {func.__name__} 必须是异步函数 (async def).")
|
||||
self.commands[name] = {
|
||||
"func": func,
|
||||
"permission": permission,
|
||||
"override_permission_check": override_permission_check,
|
||||
}
|
||||
for name in names:
|
||||
self.commands[name] = {
|
||||
"func": func,
|
||||
"permission": permission,
|
||||
"override_permission_check": override_permission_check,
|
||||
}
|
||||
return func
|
||||
return decorator
|
||||
|
||||
@@ -137,7 +138,8 @@ class MessageHandler(BaseHandler):
|
||||
permission_granted = await permission_manager.check_permission(event.user_id, permission)
|
||||
|
||||
if not permission_granted and not override_check:
|
||||
await bot.send(event, f"权限不足,需要 {permission.name} 权限")
|
||||
permission_name = permission.name if isinstance(permission, Permission) else permission
|
||||
await bot.send(event, f"权限不足,需要 {permission_name} 权限")
|
||||
return
|
||||
|
||||
await self._run_handler(
|
||||
|
||||
197
core/executor.py
197
core/executor.py
@@ -1,27 +1,184 @@
|
||||
"""
|
||||
线程池执行器
|
||||
|
||||
提供一个全局的线程池和异步接口,用于在事件循环中安全地运行同步函数。
|
||||
"""
|
||||
# -*- coding: utf-8 -*-
|
||||
import asyncio
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from functools import partial
|
||||
from typing import Any, Callable
|
||||
import docker
|
||||
from docker.tls import TLSConfig
|
||||
from typing import Dict, Any, Callable
|
||||
|
||||
# 创建一个全局的线程池,可以根据需要调整 max_workers
|
||||
executor = ThreadPoolExecutor(max_workers=10)
|
||||
from core.logger import logger
|
||||
|
||||
async def run_in_thread_pool(func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
|
||||
class CodeExecutor:
|
||||
"""
|
||||
在线程池中异步运行同步函数
|
||||
代码执行引擎,负责管理一个异步任务队列和并发的 Docker 容器执行。
|
||||
"""
|
||||
def __init__(self, bot_instance, config: Dict[str, Any]):
|
||||
"""
|
||||
初始化代码执行引擎。
|
||||
:param bot_instance: Bot 实例,用于后续的消息回复。
|
||||
:param config: 从 config.toml 加载的配置字典。
|
||||
"""
|
||||
self.bot = bot_instance
|
||||
self.task_queue = asyncio.Queue()
|
||||
|
||||
# 从传入的配置中读取 Docker 相关设置
|
||||
docker_config = config.docker
|
||||
self.docker_base_url = docker_config.get("base_url")
|
||||
self.sandbox_image = docker_config.get("sandbox_image", "python-sandbox:latest")
|
||||
self.timeout = docker_config.get("timeout", 10)
|
||||
concurrency = docker_config.get("concurrency_limit", 5)
|
||||
|
||||
self.concurrency_limit = asyncio.Semaphore(concurrency)
|
||||
self.docker_client = None
|
||||
|
||||
:param func: 要运行的同步函数
|
||||
:param args: 函数的位置参数
|
||||
:param kwargs: 函数的关键字参数
|
||||
:return: 函数的返回值
|
||||
logger.info("[CodeExecutor] 初始化 Docker 客户端...")
|
||||
try:
|
||||
if self.docker_base_url:
|
||||
# 如果配置了远程 Docker 地址,则使用 TLS 选项进行连接
|
||||
tls_config = None
|
||||
if docker_config.get("tls_verify", False):
|
||||
tls_config = TLSConfig(
|
||||
ca_cert=docker_config.get("ca_cert_path"),
|
||||
client_cert=(docker_config.get("client_cert_path"), docker_config.get("client_key_path")),
|
||||
verify=True
|
||||
)
|
||||
self.docker_client = docker.DockerClient(base_url=self.docker_base_url, tls=tls_config)
|
||||
else:
|
||||
# 否则,使用默认的本地环境连接
|
||||
self.docker_client = docker.from_env()
|
||||
|
||||
# 检查 Docker 服务是否可用
|
||||
self.docker_client.ping()
|
||||
logger.success("[CodeExecutor] Docker 客户端初始化成功,服务连接正常。")
|
||||
except docker.errors.DockerException as e:
|
||||
self.docker_client = None
|
||||
logger.error(f"无法连接到 Docker 服务,请检查 Docker 是否正在运行: {e}")
|
||||
except Exception as e:
|
||||
self.docker_client = None
|
||||
logger.error(f"初始化 Docker 客户端时发生未知错误: {e}")
|
||||
|
||||
async def add_task(self, code: str, callback: Callable[[str], asyncio.Future]):
|
||||
"""
|
||||
将代码执行任务添加到队列中。
|
||||
:param code: 待执行的 Python 代码字符串。
|
||||
:param callback: 执行完毕后用于回复结果的回调函数。
|
||||
"""
|
||||
task = {"code": code, "callback": callback}
|
||||
await self.task_queue.put(task)
|
||||
logger.info(f"[CodeExecutor] 新的代码执行任务已入队 (队列当前长度: {self.task_queue.qsize()})。")
|
||||
|
||||
async def worker(self):
|
||||
"""
|
||||
后台工作者,不断从队列中取出任务并执行。
|
||||
"""
|
||||
if not self.docker_client:
|
||||
logger.error("[CodeExecutor] Worker 无法启动,因为 Docker 客户端未初始化。")
|
||||
return
|
||||
|
||||
logger.info("[CodeExecutor] 代码执行 Worker 已启动,等待任务...")
|
||||
while True:
|
||||
task = await self.task_queue.get()
|
||||
|
||||
logger.info("[CodeExecutor] 开始处理代码执行任务。")
|
||||
|
||||
async with self.concurrency_limit:
|
||||
result_message = ""
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
# 使用 asyncio.wait_for 实现超时控制
|
||||
result_bytes = await asyncio.wait_for(
|
||||
loop.run_in_executor(
|
||||
None, # 使用默认线程池
|
||||
self._run_in_container,
|
||||
task['code']
|
||||
),
|
||||
timeout=self.timeout
|
||||
)
|
||||
|
||||
output = result_bytes.decode('utf-8').strip()
|
||||
result_message = output if output else "代码执行完毕,无输出。"
|
||||
logger.success("[CodeExecutor] 任务成功执行。")
|
||||
|
||||
except docker.errors.ImageNotFound:
|
||||
logger.error(f"[CodeExecutor] 镜像 '{self.sandbox_image}' 不存在!")
|
||||
result_message = f"执行失败:沙箱基础镜像 '{self.sandbox_image}' 不存在,请联系管理员构建。"
|
||||
except docker.errors.ContainerError as e:
|
||||
error_output = e.stderr.decode('utf-8').strip()
|
||||
result_message = f"代码执行出错:\n{error_output}"
|
||||
logger.warning(f"[CodeExecutor] 代码执行时发生错误: {error_output}")
|
||||
except docker.errors.APIError as e:
|
||||
logger.error(f"[CodeExecutor] Docker API 错误: {e}")
|
||||
result_message = "执行失败:与 Docker 服务通信时发生错误,请检查服务状态。"
|
||||
except asyncio.TimeoutError:
|
||||
result_message = f"执行超时 (超过 {self.timeout} 秒)。"
|
||||
logger.warning("[CodeExecutor] 任务执行超时。")
|
||||
except Exception as e:
|
||||
logger.exception(f"[CodeExecutor] 执行 Docker 任务时发生未知严重错误: {e}")
|
||||
result_message = "执行引擎发生内部错误,请联系管理员。"
|
||||
|
||||
# 调用回调函数回复结果
|
||||
await task['callback'](result_message)
|
||||
|
||||
self.task_queue.task_done()
|
||||
|
||||
def _run_in_container(self, code: str) -> bytes:
|
||||
"""
|
||||
同步函数:在 Docker 容器中运行代码。
|
||||
此函数通过手动管理容器生命周期来提高稳定性。
|
||||
"""
|
||||
container = None
|
||||
try:
|
||||
# 1. 创建容器
|
||||
container = self.docker_client.containers.create(
|
||||
image=self.sandbox_image,
|
||||
command=["python", "-c", code],
|
||||
mem_limit='128m',
|
||||
cpu_shares=512,
|
||||
network_disabled=True,
|
||||
log_config={'type': 'json-file', 'config': {'max-size': '1m'}},
|
||||
)
|
||||
# 2. 启动容器
|
||||
container.start()
|
||||
|
||||
# 3. 等待容器执行完成
|
||||
# 主超时由 asyncio.wait_for 控制,这里的 timeout 是一个额外的保险
|
||||
result = container.wait(timeout=self.timeout + 5)
|
||||
|
||||
# 4. 获取日志
|
||||
stdout = container.logs(stdout=True, stderr=False)
|
||||
stderr = container.logs(stdout=False, stderr=True)
|
||||
|
||||
# 5. 检查退出码,如果不为 0,则手动抛出 ContainerError
|
||||
if result.get('StatusCode', 0) != 0:
|
||||
raise docker.errors.ContainerError(
|
||||
container, result['StatusCode'], f"python -c '{code}'", self.sandbox_image, stderr
|
||||
)
|
||||
|
||||
return stdout
|
||||
|
||||
finally:
|
||||
# 6. 确保容器总是被移除
|
||||
if container:
|
||||
try:
|
||||
container.remove(force=True)
|
||||
except docker.errors.NotFound:
|
||||
# 如果容器因为某些原因已经消失,也沒关系
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.error(f"[CodeExecutor] 强制移除容器 {container.id} 时失败: {e}")
|
||||
|
||||
def initialize_executor(bot_instance, config: Dict[str, Any]):
|
||||
"""
|
||||
初始化并返回一个 CodeExecutor 实例。
|
||||
"""
|
||||
return CodeExecutor(bot_instance, config)
|
||||
|
||||
async def run_in_thread_pool(sync_func, *args, **kwargs):
|
||||
"""
|
||||
在线程池中运行同步阻塞函数,以避免阻塞 asyncio 事件循环。
|
||||
:param sync_func: 同步函数
|
||||
:param args: 位置参数
|
||||
:param kwargs: 关键字参数
|
||||
:return: 同步函数的返回值
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
# 使用 functools.partial 绑定函数和参数,以便传递给 run_in_executor
|
||||
func_to_run = partial(func, *args, **kwargs)
|
||||
# loop.run_in_executor 会返回一个 awaitable 对象
|
||||
return await loop.run_in_executor(executor, func_to_run)
|
||||
return await loop.run_in_executor(None, lambda: sync_func(*args, **kwargs))
|
||||
|
||||
@@ -227,6 +227,14 @@ class PermissionManager:
|
||||
Returns:
|
||||
bool: 如果用户权限 >= 所需权限,返回 True,否则返回 False
|
||||
"""
|
||||
# 如果传入的是字符串,先转换为 Permission 对象
|
||||
if isinstance(required_permission, str):
|
||||
required_permission = _PERMISSIONS.get(required_permission.lower())
|
||||
if not required_permission:
|
||||
# 如果是无效的权限字符串,默认拒绝
|
||||
logger.warning(f"检测到无效的权限检查字符串: {required_permission}")
|
||||
return False
|
||||
|
||||
user_permission = await self.get_user_permission(user_id)
|
||||
return user_permission >= required_permission
|
||||
|
||||
@@ -249,4 +257,21 @@ class PermissionManager:
|
||||
|
||||
|
||||
# 全局权限管理器实例
|
||||
permission_manager = PermissionManager()
|
||||
permission_manager = PermissionManager()
|
||||
|
||||
def require_admin(func):
|
||||
"""
|
||||
一个装饰器,用于限制命令只能由管理员执行。
|
||||
"""
|
||||
from functools import wraps
|
||||
from models.events.message import MessageEvent
|
||||
|
||||
@wraps(func)
|
||||
async def wrapper(event: MessageEvent, *args, **kwargs):
|
||||
user_id = event.user_id
|
||||
if await permission_manager.check_permission(user_id, ADMIN):
|
||||
return await func(event, *args, **kwargs)
|
||||
else:
|
||||
await event.reply("抱歉,您没有权限执行此命令。")
|
||||
return None
|
||||
return wrapper
|
||||
|
||||
Reference in New Issue
Block a user