shell-command-runner

import { tool, Tool, ToolsProviderController } from "@lmstudio/sdk";
import { exec, ExecOptions } from "child_process";
import { z } from "zod";
import path from "path";
import { configSchematics } from "./config";

interface ExecResult {
  stdout: string;
  stderr: string;
  error: Error | null;
  timedOut: boolean;
}

const runShellCommand = (command: string, options: ExecOptions, timeoutMs: number): Promise<ExecResult> => {
  return new Promise((resolve) => {
    // NODEJS EXEC WRAPPER
    const child = exec(command, { ...options, timeout: timeoutMs }, (error, stdout, stderr) => {
      resolve({
        stdout: stdout.trim(),
        stderr: stderr.trim(),
        error,
        timedOut: false,
      });
    });

    child.on('error', (err: any) => {
      if (err.code === 'ETIMEDOUT') {
        resolve({ stdout: '', stderr: 'Command timed out.', error: err, timedOut: true });
      }
    });
  });
};

export async function toolsProvider(ctl: ToolsProviderController): Promise<Tool[]> {
  const config = ctl.getPluginConfig(configSchematics);
  const osType = config.get("operatingSystem");
  
  // --- 1. DETERMINE SHELL EXECUTABLE ---
  // Reverted to CMD.exe for Windows to preserve native PATH behavior (fixes 'cargo' not found)
  const shellExecutable = osType === "windows" ? "cmd.exe" : "/bin/bash";

  // --- 2. OS HINTS ---
  let osHints = "";
  if (osType === "windows") {
    osHints = `
    TARGET SHELL: WINDOWS CMD (Command Prompt).
    - Standard Syntax: Use 'dir', 'type', 'copy', 'cargo', 'python'.
    - POWERSHELL: This shell does NOT support PowerShell syntax (ls, Get-ChildItem) directly.
      IF YOU NEED POWERSHELL: You must prefix the command like this:
      powershell -Command "Get-ChildItem"
    - Chaining: Use '&&' (not ';').
    `;
  } else {
    osHints = `
    TARGET SHELL: LINUX/MACOS (Bash).
    - Usage: Standard bash commands.
    - List: 'ls -la'
    `;
  }

  // --- 3. SANDBOX SETUP ---
  const sandboxRootStr = config.get("homeDirectory");
  const sandboxRoot = sandboxRootStr && sandboxRootStr.trim() !== "" 
    ? path.resolve(sandboxRootStr) 
    : null; 

  let sandboxWarning = "";
  if (sandboxRoot) {
    sandboxWarning = `
    SECURITY RESTRICTION ACTIVE:
    - Working Directory: "${sandboxRoot}"
    - You CANNOT access files outside this folder.
    - DO NOT use absolute paths.
    `;
  }

  const dynamicDescription = `
  FALLBACK TOOL: Executes a command in the system shell.
  PRIORITY: Check for specialized tools (like 'read_file') BEFORE using this.
  ${sandboxWarning}
  ${osHints}
  `;

  const executeCommandTool = tool({
    name: "run_shell_command", 
    description: dynamicDescription,
    parameters: {
      command: z.string().describe(`The command to run.`),
      cwd: z.string().optional().describe("Subdirectory to run in. Must be inside the Sandbox."),
      timeout: z.number().optional().describe("Execution timeout in milliseconds."),
    },
    implementation: async ({ command, cwd, timeout }, { status, warn }) => {
      const currentConfig = ctl.getPluginConfig(configSchematics);
      const allowAuto = currentConfig.get("allowAutoExecution");
      
      const policy = currentConfig.get("executionPolicy"); // "allow_all" or "allow_only"
      const allowedStr = currentConfig.get("allowedCommands");
      const forbiddenStr = currentConfig.get("forbiddenCommands");
      const extraPathsStr = currentConfig.get("additionalSearchPaths");

      // --- CHECK 1: AUTO EXECUTION ---
      if (!allowAuto) {
        warn(`Blocked command: "${command}"`);
        return `PERMISSION DENIED: 'Allow Automatic Execution' is disabled.`;
      }

      // --- CHECK 2: BLACKLIST ---
      if (forbiddenStr && forbiddenStr.trim().length > 0) {
        const forbiddenList = forbiddenStr.split(",").map(s => s.trim().toLowerCase());
        const commandLower = command.toLowerCase();
        const isForbidden = forbiddenList.some(badCmd => {
             return commandLower.startsWith(badCmd + " ") || 
                    commandLower === badCmd ||
                    commandLower.includes(" " + badCmd + " "); 
        });
        if (isForbidden) return `SECURITY ERROR: The command '${command}' is forbidden by the Blacklist.`;
      }

      // --- CHECK 3: WHITELIST ---
      if (policy === "allow_only") {
        if (!allowedStr || allowedStr.trim() === "") {
            return `SECURITY ERROR: Policy is 'Allow Only' but the Allowed Commands list is empty.`;
        }
        const allowedList = allowedStr.split(",").map(s => s.trim().toLowerCase());
        const commandLower = command.toLowerCase();
        
        // For whitelist, we check if the command starts with an allowed token.
        // We also need to handle "powershell -Command" cases specially if we want to allow them.
        const isAllowed = allowedList.some(allowedCmd => {
            // Check matches
            const match = commandLower === allowedCmd || commandLower.startsWith(allowedCmd + " ");
            // Special case: If user whitelisted "powershell", allow "powershell -Command ..."
            return match;
        });

        if (!isAllowed) {
            return `SECURITY ERROR: The command '${command}' is NOT in the Allowed Commands list.`;
        }
      }

      // --- CHECK 4: SANDBOX ---
      const effectiveRoot = sandboxRoot || process.cwd();
      let targetDir = effectiveRoot;
      if (cwd) targetDir = path.resolve(effectiveRoot, cwd);

      if (sandboxRoot) {
        const relativePath = path.relative(sandboxRoot, targetDir);
        if (relativePath.startsWith('..') || path.isAbsolute(relativePath)) {
           return `SECURITY ERROR: Access denied. Target '${targetDir}' is outside sandbox '${sandboxRoot}'.`;
        }
        if (/[a-zA-Z]:/.test(command)) return `SECURITY ERROR: Absolute paths (Drive Letters) are forbidden while sandboxed.`;
        if (command.includes("..")) return `SECURITY ERROR: Parent directory traversal ('..') is forbidden while sandboxed.`;
      }

      // --- SETUP ENVIRONMENT ---
      // We still include this helper, just in case CMD also needs help finding path variables.
      const env = { ...process.env };
      if (extraPathsStr && extraPathsStr.trim() !== "") {
        const extraPaths = extraPathsStr.split(",").map(p => p.trim()).join(path.delimiter);
        const pathKey = process.platform === 'win32' ? 'Path' : 'PATH'; 
        const actualPathKey = Object.keys(env).find(k => k.toLowerCase() === 'path') || pathKey;
        env[actualPathKey] = `${env[actualPathKey] || ''}${path.delimiter}${extraPaths}`;
      }

      status(`Running: ${command.slice(0, 50)}...`);

      try {
        const timeLimit = timeout || currentConfig.get("defaultTimeout");
        
        const result = await runShellCommand(command, { 
            cwd: targetDir,
            shell: shellExecutable, // Back to cmd.exe
            env: env 
        }, timeLimit);

        let response = "";
        
        if (result.stdout) {
          const maxLen = 4000;
          const output = result.stdout.length > maxLen 
            ? result.stdout.slice(0, maxLen) + "\n...[Output Truncated]" 
            : result.stdout;
          response += `STDOUT:\n${output}\n`;
        }

        if (result.stderr) {
          response += `STDERR:\n${result.stderr}\n`;
        }

        if (result.error) {
           response += `EXIT_CODE: ${"code" in result.error ? result.error.code : "Unknown"}\n`;
           if (!result.stderr) response += `ERROR: ${result.error.message}\n`;
        } else {
           if (!result.stdout && !result.stderr) response = "Command executed successfully.";
        }

        return response.trim();

      } catch (err: any) {
        return `PLUGIN EXCEPTION: ${err.message}`;
      }
    },
  });

  return [executeCommandTool];
}