1  import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
  2  import { z } from "zod";
  3  import type { EasyShellClient } from "../client.js";
  4  import { toMcpError } from "../errors.js";
  5  
  6  type ExecuteScriptParams = {
  7    name: string;
  8    agentIds?: string[];
  9    clusterIds?: number[];
 10    tagIds?: number[];
 11    scriptId?: number;
 12    scriptContent?: string;
 13    timeoutSeconds?: number;
 14  };
 15  
 16  function isRecord(value: unknown): value is Record<string, unknown> {
 17    return typeof value === "object" && value !== null;
 18  }
 19  
 20  function hasRiskApprovalCode(value: unknown): boolean {
 21    if (!isRecord(value)) {
 22      return false;
 23    }
 24  
 25    const code = value["code"];
 26    return code === 449;
 27  }
 28  
 29  export function registerScriptTools(server: McpServer, client: EasyShellClient) {
 30    server.tool(
 31      "list_scripts",
 32      "List all scripts in the script library. Includes templates and user-created scripts.",
 33      {},
 34      async () => {
 35        try {
 36          const result = await client.get("/api/v1/script/list");
 37          return {
 38            content: [{ type: "text" as const, text: JSON.stringify(result, null, 2) }],
 39          };
 40        } catch (error) {
 41          return toMcpError(error);
 42        }
 43      }
 44    );
 45  
 46    server.tool(
 47      "create_script",
 48      "Create a new script in the script library.",
 49      {
 50        name: z.string(),
 51        content: z.string(),
 52        description: z.string().optional(),
 53        scriptType: z.enum(["shell", "python", "perl"]).default("shell").optional(),
 54      },
 55      async ({
 56        name,
 57        content,
 58        description,
 59        scriptType,
 60      }: {
 61        name: string;
 62        content: string;
 63        description?: string;
 64        scriptType?: "shell" | "python" | "perl";
 65      }) => {
 66        try {
 67          const result = await client.post("/api/v1/script", {
 68            name,
 69            content,
 70            description,
 71            scriptType,
 72          });
 73  
 74          return {
 75            content: [{ type: "text" as const, text: JSON.stringify(result, null, 2) }],
 76          };
 77        } catch (error) {
 78          return toMcpError(error);
 79        }
 80      }
 81    );
 82  
 83    server.tool(
 84      "execute_script",
 85      "Execute a script on one or more hosts. Returns a task ID that can be used with get_task_detail to check results. Can specify either scriptId (from library) or scriptContent (inline).",
 86      {
 87        name: z.string().describe("Task name"),
 88        agentIds: z.array(z.string()).optional().describe("Host IDs to execute on"),
 89        clusterIds: z.array(z.number()).optional().describe("Cluster IDs"),
 90        tagIds: z
 91          .array(z.number())
 92          .optional()
 93          .describe("Tag IDs — execute on all hosts with these tags"),
 94        scriptId: z.number().optional().describe("Script ID from library"),
 95        scriptContent: z.string().optional().describe("Inline script content"),
 96        timeoutSeconds: z.number().default(3600).optional(),
 97      },
 98      async ({
 99        name,
100        agentIds,
101        clusterIds,
102        tagIds,
103        scriptId,
104        scriptContent,
105        timeoutSeconds,
106      }: ExecuteScriptParams) => {
107        try {
108          const result = await client.post("/api/v1/task", {
109            name,
110            agentIds,
111            clusterIds,
112            tagIds,
113            scriptId,
114            scriptContent,
115            timeoutSeconds,
116          });
117  
118          return {
119            content: [{ type: "text" as const, text: JSON.stringify(result, null, 2) }],
120          };
121        } catch (error) {
122          if (hasRiskApprovalCode(error)) {
123            return {
124              content: [{ type: "text" as const, text: JSON.stringify(error, null, 2) }],
125            };
126          }
127          return toMcpError(error);
128        }
129      }
130    );
131  }