feat: 添加Docker沙箱代码执行功能
- 新增Docker沙箱执行环境，提供安全隔离的代码执行能力 - 重构code_py插件，使用Docker容器替代子进程执行 - 添加docker配置项和权限检查功能 - 实现代码执行队列和并发控制 - 新增广播插件，仅限管理员使用
This commit is contained in:
@@ -1,176 +1,156 @@
|
||||
"""
|
||||
code_py插件
|
||||
|
||||
输入/code py回车再加上python代码，机器人就会执行代码并返回执行结果。
|
||||
"""
|
||||
|
||||
# -*- coding: utf-8 -*-
|
||||
import html
|
||||
import textwrap
|
||||
import asyncio
|
||||
import re
|
||||
import sys
|
||||
import tempfile
|
||||
import os
|
||||
from typing import Tuple, Set
|
||||
from typing import Dict
|
||||
|
||||
from core.bot import Bot
|
||||
from core.command_manager import matcher
|
||||
from core.executor import run_in_thread_pool
|
||||
from models import MessageEvent
|
||||
from core.permission_manager import ADMIN
|
||||
from core.logger import logger
|
||||
|
||||
__plugin_meta__ = {
|
||||
"name": "code_py",
|
||||
"description": "提供执行python代码的功能",
|
||||
"usage": "/code_py - 进入交互模式，等待输入代码块\n/code_py [单行代码] - 快速执行单行代码",
|
||||
"name": "Python 代码执行",
|
||||
"description": "在安全的沙箱环境中执行 Python 代码片段，支持单行、多行和转发回复。",
|
||||
"usage": "/py <单行代码>\n/code_py <单行代码>\n/py (进入多行输入模式)",
|
||||
}
|
||||
|
||||
# --- 安全配置：危险模块和内置函数黑名单 ---
|
||||
DANGEROUS_MODULES = [
|
||||
"os", "sys", "subprocess", "shutil", "socket", "requests", "urllib",
|
||||
"http", "ftplib", "telnetlib", "ctypes", "_thread", "multiprocessing",
|
||||
"asyncio",
|
||||
]
|
||||
DANGEROUS_BUILTINS = [
|
||||
"__import__", "open", "exec", "eval", "compile", "input", "breakpoint"
|
||||
]
|
||||
# --- 会话状态管理 ---
|
||||
# 结构: {(user_id, group_id): asyncio.TimerHandle}
|
||||
multi_line_sessions: Dict[tuple, asyncio.TimerHandle] = {}
|
||||
|
||||
# 编译后的正则表达式，用于分割语句
|
||||
STATEMENT_SPLIT_PATTERN = re.compile(r'[;\n]')
|
||||
# 编译后的正则表达式，用于查找危险的内置函数调用
|
||||
BUILTIN_CALL_PATTERN = re.compile(r'\b(' + '|'.join(DANGEROUS_BUILTINS) + r')\s*\(')
|
||||
|
||||
def is_code_safe(code: str) -> Tuple[bool, str]:
|
||||
async def reply_as_forward(event: MessageEvent, input_code: str, output_result: str):
|
||||
"""
|
||||
检查代码中是否包含危险的模块导入或内置函数调用。
|
||||
将输入和输出打包成转发消息进行回复。
|
||||
参考 forward_test.py 的实现，兼容私聊和群聊。
|
||||
"""
|
||||
# 1. 检查危险的内置函数
|
||||
found_builtins = BUILTIN_CALL_PATTERN.search(code)
|
||||
if found_builtins:
|
||||
return False, f"检测到不允许的内置函数调用：'{found_builtins.group(1)}'"
|
||||
|
||||
# 2. 检查危险的模块导入
|
||||
statements = STATEMENT_SPLIT_PATTERN.split(code)
|
||||
for statement in statements:
|
||||
statement = statement.strip()
|
||||
if not statement:
|
||||
continue
|
||||
parts = statement.split()
|
||||
if not parts:
|
||||
continue
|
||||
if parts[0] == 'from' and len(parts) > 1:
|
||||
module_name = parts[1].strip()
|
||||
if module_name in DANGEROUS_MODULES:
|
||||
return False, f"检测到不允许的模块导入：'{module_name}'"
|
||||
elif parts[0] == 'import' and len(parts) > 1:
|
||||
modules_str = ' '.join(parts[1:])
|
||||
imported_modules = [m.strip() for m in modules_str.split(',')]
|
||||
for module_name in imported_modules:
|
||||
actual_module_name = module_name.split()[0]
|
||||
if actual_module_name in DANGEROUS_MODULES:
|
||||
return False, f"检测到不允许的模块导入：'{actual_module_name}'"
|
||||
return True, ""
|
||||
|
||||
async def run_code_in_subprocess(code_str: str, timeout: float = 10.0) -> Tuple[str, str]:
|
||||
"""
|
||||
在子进程中安全地执行Python代码。
|
||||
"""
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False, encoding="utf-8") as tf:
|
||||
tf.write(code_str)
|
||||
tf_path = tf.name
|
||||
try:
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
sys.executable, tf_path,
|
||||
stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
|
||||
)
|
||||
try:
|
||||
out_bytes, err_bytes = await asyncio.wait_for(proc.communicate(), timeout=timeout)
|
||||
except asyncio.TimeoutError:
|
||||
proc.kill()
|
||||
await proc.communicate()
|
||||
return "", f"执行超时（>{timeout}s）"
|
||||
return out_bytes.decode(errors="ignore"), err_bytes.decode(errors="ignore")
|
||||
finally:
|
||||
try:
|
||||
os.remove(tf_path)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def process_and_reply(bot: Bot, event: MessageEvent, code: str):
|
||||
"""
|
||||
核心处理逻辑：安全检查、执行代码并回复结果。
|
||||
"""
|
||||
safe, message = await run_in_thread_pool(is_code_safe, code)
|
||||
if not safe:
|
||||
await event.reply(f"代码安全检查未通过：\n{message}")
|
||||
return
|
||||
|
||||
try:
|
||||
stdout, stderr = await run_code_in_subprocess(code, timeout=10.0)
|
||||
except Exception as e:
|
||||
await event.reply(f"执行失败：{e}")
|
||||
return
|
||||
|
||||
resp = stderr.strip() or stdout.strip() or "(无输出)"
|
||||
MAX = 1500
|
||||
if len(resp) > MAX:
|
||||
resp = resp[:MAX] + "\n...输出被截断..."
|
||||
|
||||
bot = event.bot
|
||||
|
||||
# 1. 构建消息节点列表
|
||||
nodes = [
|
||||
bot.build_forward_node(user_id=event.self_id, nickname="输入代码", message=code),
|
||||
bot.build_forward_node(user_id=event.self_id, nickname="执行结果", message=resp),
|
||||
bot.build_forward_node(
|
||||
user_id=event.user_id,
|
||||
nickname=event.sender.nickname or str(event.user_id),
|
||||
message=f"--- Your Code ---\n{input_code}"
|
||||
),
|
||||
bot.build_forward_node(
|
||||
user_id=event.self_id,
|
||||
nickname="Code Executor",
|
||||
message=f"--- Execution Result ---\n{output_result}"
|
||||
)
|
||||
]
|
||||
|
||||
try:
|
||||
# 2. 发送合并转发消息
|
||||
await bot.send_forwarded_messages(event, nodes)
|
||||
except Exception as e:
|
||||
await event.reply(f"结果发送失败: {e}\n\n{resp}")
|
||||
logger.error(f"[code_py] 发送转发消息失败: {e}")
|
||||
# 降级为普通消息回复
|
||||
await event.reply(f"--- 你的代码 ---\n{input_code}\n--- 执行结果 ---\n{output_result}")
|
||||
|
||||
# --- 交互式会话状态 ---
|
||||
# 使用集合存储正在等待代码输入的用户标识
|
||||
waiting_users: Set[str] = set()
|
||||
|
||||
def get_session_id(event: MessageEvent) -> str:
|
||||
"""根据事件类型生成唯一的会话ID"""
|
||||
if hasattr(event, 'group_id'):
|
||||
# 群聊会话ID
|
||||
return f"group_{event.group_id}-{event.user_id}"
|
||||
else:
|
||||
# 私聊会话ID
|
||||
return f"private_{event.user_id}"
|
||||
|
||||
@matcher.command("code_py")
|
||||
async def handle_code_command(bot: Bot, event: MessageEvent, args: list[str]):
|
||||
# 模式一：快速执行单行代码
|
||||
if args:
|
||||
code = " ".join(args)
|
||||
await process_and_reply(bot, event, code)
|
||||
async def execute_code(event: MessageEvent, code: str):
|
||||
"""
|
||||
核心代码执行逻辑。
|
||||
"""
|
||||
code_executor = getattr(event.bot, 'code_executor', None)
|
||||
if not code_executor or not code_executor.docker_client:
|
||||
await event.reply("代码执行服务当前不可用，请检查 Docker 连接配置。")
|
||||
return
|
||||
|
||||
# 模式二：进入交互模式
|
||||
session_id = get_session_id(event)
|
||||
if session_id in waiting_users:
|
||||
await event.reply("您已经有一个正在等待输入的code会话了，请直接发送代码。")
|
||||
return
|
||||
# 修改 add_task，让它能直接接收回复函数
|
||||
await code_executor.add_task(
|
||||
code,
|
||||
lambda result: reply_as_forward(event, code, result)
|
||||
)
|
||||
await event.reply("代码已提交至沙箱执行队列，请稍候...")
|
||||
|
||||
def cleanup_session(session_key: tuple):
|
||||
"""
|
||||
清理超时的会话。
|
||||
"""
|
||||
if session_key in multi_line_sessions:
|
||||
del multi_line_sessions[session_key]
|
||||
logger.info(f"[code_py] 会话 {session_key} 已超时，自动取消。")
|
||||
|
||||
def normalize_code(code: str) -> str:
|
||||
"""
|
||||
规范化用户输入的 Python 代码字符串。
|
||||
|
||||
主要处理两个问题：
|
||||
1. 对消息中可能存在的 HTML 实体进行解码 (e.g., [ -> [)。
|
||||
2. 移除整个代码块的公共前导缩进，以修复因复制粘贴导致的多余缩进。
|
||||
|
||||
:param code: 原始代码字符串。
|
||||
:return: 规范化后的代码字符串。
|
||||
"""
|
||||
# 1. 解码 HTML 实体
|
||||
code = html.unescape(code)
|
||||
|
||||
# 2. 移除公共前导缩进
|
||||
try:
|
||||
code = textwrap.dedent(code)
|
||||
except Exception:
|
||||
# 在某些情况下（例如，不一致的缩进），dedent 可能会失败，
|
||||
# 但我们不希望因此中断流程，所以捕获异常并继续。
|
||||
pass
|
||||
|
||||
waiting_users.add(session_id)
|
||||
await event.reply("请在下一条消息中发送要执行的Python代码块。（发送“取消”可退出）")
|
||||
return code.strip()
|
||||
|
||||
|
||||
@matcher.command("py", "python", "code_py", permission=ADMIN)
|
||||
async def code_py_main(event: MessageEvent, args: list[str]):
|
||||
"""
|
||||
/py 命令的主入口。
|
||||
- 如果有参数，直接执行。
|
||||
- 如果没有参数，开启多行输入模式。
|
||||
"""
|
||||
code_to_run = " ".join(args)
|
||||
|
||||
if code_to_run:
|
||||
# 单行模式，对代码进行规范化处理
|
||||
normalized_code = normalize_code(code_to_run)
|
||||
if not normalized_code:
|
||||
await event.reply("代码为空或格式错误，请输入有效的代码。")
|
||||
return
|
||||
await execute_code(event, normalized_code)
|
||||
else:
|
||||
# 多行模式
|
||||
# 使用 getattr 兼容私聊和群聊
|
||||
session_key = (event.user_id, getattr(event, 'group_id', 'private'))
|
||||
|
||||
# 如果上一个会话的超时任务还在，先取消它
|
||||
if session_key in multi_line_sessions:
|
||||
multi_line_sessions[session_key].cancel()
|
||||
|
||||
await event.reply("已进入多行代码输入模式，请直接发送你的代码。\n(60秒内无操作将自动取消)")
|
||||
|
||||
# 设置 60 秒超时
|
||||
loop = asyncio.get_running_loop()
|
||||
timeout_handler = loop.call_later(
|
||||
60,
|
||||
cleanup_session,
|
||||
session_key
|
||||
)
|
||||
multi_line_sessions[session_key] = timeout_handler
|
||||
|
||||
@matcher.on_message()
|
||||
async def handle_code_input(bot: Bot, event: MessageEvent):
|
||||
session_id = get_session_id(event)
|
||||
|
||||
# 检查用户是否处于等待状态
|
||||
if session_id in waiting_users:
|
||||
# 从等待集合中移除，无论输入是什么
|
||||
waiting_users.remove(session_id)
|
||||
async def handle_multi_line_code(event: MessageEvent):
|
||||
"""
|
||||
通用消息处理器，用于捕获多行模式下的代码输入。
|
||||
"""
|
||||
# 使用 getattr 兼容私聊和群聊
|
||||
session_key = (event.user_id, getattr(event, 'group_id', 'private'))
|
||||
if session_key in multi_line_sessions:
|
||||
# 取消超时任务
|
||||
multi_line_sessions[session_key].cancel()
|
||||
del multi_line_sessions[session_key]
|
||||
|
||||
# 对多行代码进行规范化处理
|
||||
normalized_code = normalize_code(event.raw_message)
|
||||
|
||||
# 处理取消操作
|
||||
if event.raw_message.strip() == "取消":
|
||||
await event.reply("已取消输入。")
|
||||
return True # 消费事件
|
||||
|
||||
# 执行代码
|
||||
await process_and_reply(bot, event, event.raw_message)
|
||||
return True # 消费事件，防止被其他指令匹配
|
||||
|
||||
# 如果用户不在等待状态，则不处理
|
||||
return False
|
||||
|
||||
if not normalized_code:
|
||||
await event.reply("捕获到的代码为空或格式错误，已取消输入。")
|
||||
return
|
||||
|
||||
await execute_code(event, normalized_code)
|
||||
return True # 消费事件，防止其他处理器响应
|
||||
|
||||
Reference in New Issue
Block a user