115 lines
4.6 KiB
Python
115 lines
4.6 KiB
Python
# system_tools.py
|
|
import os
|
|
import asyncssh
|
|
from mcp.server.fastmcp import FastMCP, Context
|
|
from mcp.server.session import ServerSession
|
|
from lifespan import SharedAppContext # 导入共享上下文的类型定义
|
|
|
|
|
|
def create_system_mcp() -> FastMCP:
|
|
"""创建一个包含系统操作工具的 MCP 实例。"""
|
|
|
|
system_mcp = FastMCP(
|
|
name="System Tools",
|
|
instructions="用于在远程服务器上进行基本文件和目录操作的工具集。",
|
|
# 将 MCP 实例的 HTTP 路径设置为根，以便挂载在 /system 下 [1]
|
|
streamable_http_path="/"
|
|
)
|
|
|
|
@system_mcp.tool()
|
|
async def list_files(ctx: Context[ServerSession, dict], path: str = ".") -> str:
|
|
"""
|
|
列出远程服务器安全沙箱路径下的文件和目录。
|
|
Args:
|
|
path: 要查看的相对路径，默认为当前沙箱目录。
|
|
"""
|
|
try:
|
|
# 从上下文中获取共享的 SSH 连接和沙箱路径 [1]
|
|
shared_ctx: SharedAppContext = ctx.request_context.lifespan_context["shared_context"]
|
|
|
|
conn = shared_ctx.ssh_connection
|
|
sandbox_path = shared_ctx.sandbox_path
|
|
|
|
# 安全检查：防止目录穿越攻击
|
|
if ".." in path.split('/'):
|
|
return "错误: 不允许使用 '..' 访问上级目录。"
|
|
|
|
# 安全地拼接路径
|
|
target_path = os.path.join(sandbox_path, path)
|
|
await ctx.info(f"正在列出安全路径: {target_path}")
|
|
|
|
# 使用 conn.run() 执行 ls -l 命令
|
|
command_to_run = f'ls -la {conn.escape(target_path)}'
|
|
result = await conn.run(command_to_run, check=True)
|
|
return result.stdout.strip()
|
|
except asyncssh.ProcessError as e:
|
|
error_message = f"命令执行失败: {e.stderr.strip()}"
|
|
await ctx.error(error_message)
|
|
return error_message
|
|
except Exception as e:
|
|
error_message = f"执行 ls 命令时发生未知错误: {e}"
|
|
await ctx.error(error_message)
|
|
return error_message
|
|
|
|
@system_mcp.tool()
|
|
async def read_file(ctx: Context[ServerSession, dict], file_path: str) -> str:
|
|
"""
|
|
读取远程服务器安全沙箱内指定文件的内容。
|
|
Args:
|
|
file_path: 相对于沙箱目录的文件路径。
|
|
"""
|
|
try:
|
|
shared_ctx: SharedAppContext = ctx.request_context.lifespan_context["shared_context"]
|
|
|
|
conn = shared_ctx.ssh_connection
|
|
sandbox_path = shared_ctx.sandbox_path
|
|
# 安全检查：防止目录穿越
|
|
if ".." in file_path.split('/'):
|
|
return "错误: 不允许使用 '..' 访问上级目录。"
|
|
|
|
target_path = os.path.join(sandbox_path, file_path)
|
|
await ctx.info(f"正在读取远程文件: {target_path}")
|
|
|
|
async with conn.start_sftp_client() as sftp:
|
|
async with sftp.open(target_path, 'r') as f:
|
|
content = await f.read()
|
|
return content
|
|
except FileNotFoundError:
|
|
error_message = f"读取文件失败: 文件 '{file_path}' 不存在。"
|
|
await ctx.error(error_message)
|
|
return error_message
|
|
except Exception as e:
|
|
await ctx.error(f"读取文件时发生未知错误: {e}")
|
|
return f"错误: {e}"
|
|
|
|
@system_mcp.tool()
|
|
async def write_file(ctx: Context[ServerSession, dict], file_path: str, content: str) -> str:
|
|
"""
|
|
向远程服务器安全沙箱内的文件写入内容。如果文件已存在，则会覆盖它。
|
|
Args:
|
|
file_path: 相对于沙箱目录的文件路径。
|
|
content: 要写入文件的文本内容。
|
|
"""
|
|
try:
|
|
shared_ctx: SharedAppContext = ctx.request_context.lifespan_context["shared_context"]
|
|
|
|
conn = shared_ctx.ssh_connection
|
|
sandbox_path = shared_ctx.sandbox_path
|
|
|
|
# 安全检查
|
|
if ".." in file_path.split('/'):
|
|
return "错误: 不允许使用 '..' 访问上级目录。"
|
|
|
|
target_path = os.path.join(sandbox_path, file_path)
|
|
await ctx.info(f"正在向远程文件写入: {target_path}")
|
|
|
|
async with conn.start_sftp_client() as sftp:
|
|
async with sftp.open(target_path, 'w') as f:
|
|
await f.write(content)
|
|
return f"内容已成功写入到沙箱文件 '{file_path}'。"
|
|
except Exception as e:
|
|
await ctx.error(f"写入文件时发生未知错误: {e}")
|
|
return f"错误: {e}"
|
|
|
|
return system_mcp
|