# server.py

import asyncio
import asyncssh  # 用于 SSH 连接的库
from mcp.server.fastmcp import FastMCP, Context

# --- 预设的服务器连接信息 ---
# 警告：在生产环境中，不要将密钥和密码硬编码在代码里。
# 最好使用环境变量、配置文件或专门的密钥管理服务。
REMOTE_HOST = '202.121.182.208'  # 替换为你的服务器地址
REMOTE_USER = 'koko125'  # 替换为你的用户名
PRIVATE_KEY_PATH = 'D:/tool/tool/id_rsa.txt'  # 替换为你的私钥文件路径

# 创建一个 MCP 服务器实例
mcp = FastMCP("远程服务器工具")


@mcp.tool()
async def execute_remote_command(command: str, ctx: Context) -> str:
    """
    在远程服务器上执行一个shell命令并返回其输出。

    Args:
        command: 要在远程服务器上执行的命令字符串。
    """
    await ctx.info(f"准备在 {REMOTE_HOST} 上执行命令: '{command}'")

    try:
        # 建立 SSH 连接
        async with asyncssh.connect(REMOTE_HOST, username=REMOTE_USER, client_keys=[PRIVATE_KEY_PATH]) as conn:
            # 执行命令
            result = await conn.run(command, check=True)

            # 成功执行，返回标准输出
            output = result.stdout.strip()
            await ctx.info(f"命令成功执行，返回输出。")
            await ctx.debug(f"输出内容: {output}")
            return output

    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"发生未知错误: {str(e)}"
        await ctx.error(error_message)
        return error_message


# 这部分使得你可以直接运行 `python server.py`
if __name__ == "__main__":
    # mcp.run() # 这是默认的 stdio 模式
    mcp.run(transport="streamable-http") # 改为 streamable-http 模式