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.
  
  TEST MODE: If 'Test Mode' is enabled in plugin settings, this tool only checks if the command is allowed/forbidden and returns status without executing.
  ${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");
      const testMode = currentConfig.get("testMode");

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

      // --- SETUP ENVIRONMENT (needed for cd exception handling) ---
      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}`;
      }

      // --- EXCEPTION: ALLOW cd {sandboxRoot} Bypassing Both Blacklist and Whitelist ---
      const isCdCommand = command.trim().toLowerCase().startsWith("cd ");
      if (isCdCommand && sandboxRoot) {
        // Extract the path portion - stop at && or other shell operators
        const afterCd = command.substring(3).trim();
        
        // Find where the path ends (before && or end of string)
        let pathEnd = afterCd.length;
        const andIndex = afterCd.indexOf("&&");
        if (andIndex !== -1) {
          pathEnd = andIndex;
        }
        
        const cdTargetRaw = afterCd.substring(0, pathEnd).trim();
        const resolvedCdTarget = path.resolve(sandboxRoot, cdTargetRaw);
        
        // Check if this is "cd {sandboxRoot}" or "cd {sandboxRoot} && ..."
        const remainingAfterTarget = afterCd.substring(pathEnd).trim();
        
        const isValidCdPattern = resolvedCdTarget === sandboxRoot && 
                                  (remainingAfterTarget === "" || remainingAfterTarget.startsWith("&&"));
        
        if (isValidCdPattern) {
          // This is 'cd {sandboxRoot}' or 'cd {sandboxRoot} && ...' - allow it unconditionally
          status(`Running: ${command.slice(0, 50)}...`);
          
          try {
            const timeLimit = timeout || currentConfig.get("defaultTimeout");
            
            const result = await runShellCommand(command, { 
                cwd: sandboxRoot,
                shell: shellExecutable,
                env: env 
            }, timeLimit);

            let response = "";
            let exitCode = null;
            
            const hasError = result.error && result.error.code !== undefined;
            
            if (hasError) {
               exitCode = "code" in result.error ? result.error.code : null;
               
               if (exitCode !== 1) {
                 response += `ERROR: ${result.error.message}\n`;
               }
               
               if (result.stderr && !result.error.message.includes(result.stderr)) {
                 response += `STDERR:\n${result.stderr}\n`;
               }
            } else if (!result.stdout && !result.stderr) {
               response = "Command executed successfully.";
            }
            
            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 (hasError) {
               response += `EXIT_CODE: ${exitCode !== null ? exitCode : "Unknown"}\n`;
            }

            return response.trim();

          } catch (err: any) {
            return `PLUGIN EXCEPTION: ${err.message}`;
          }
        } else {
          // cd command but not to sandbox root - blocked by sandbox restriction
          return `SECURITY ERROR: 'cd' commands are only allowed to the sandbox root '${sandboxRoot}'. Target '${resolvedCdTarget}' is outside the sandbox.`;
        }
      }

      // --- 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}'.`;
        }
        // Only block Windows absolute paths (e.g., C:\path\to\file or D:/path/to/file)
        // Match drive letter at start of command or after whitespace/semicolon, followed by colon and path separator
        if (/^\s*[a-zA-Z]:[\\/]/.test(command) || /;\s*[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.`;
      }

      // --- TEST MODE CHECK (AFTER ALL SECURITY CHECKS PASS) ---
      if (testMode) {
        return `COMMAND ALLOWED: The command '${command}' passes all security checks.`;
      }

      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 = "";
        let exitCode = null;
        
        // Only report error if there's an actual execution error, not just non-zero exit code
        // Node.js exec() passes error for non-zero exit codes, but that's normal (e.g., grep returns 1 when no matches)
        const hasError = result.error && result.error.code !== undefined;
        
        if (hasError) {
           exitCode = "code" in result.error ? result.error.code : null;
           
           // Only show ERROR message if exit code is not 1 (grep no matches is normal behavior)
           // Exit code 1 from grep with no matches should just show EXIT_CODE, not ERROR
           if (exitCode !== 1) {
             response += `ERROR: ${result.error.message}\n`;
           }
           
           // Only show STDERR separately if it's different from what's in error.message
           // (error.message already includes stderr for command failures)
           if (result.stderr && !result.error.message.includes(result.stderr)) {
             response += `STDERR:\n${result.stderr}\n`;
           }
        } else if (!result.stdout && !result.stderr) {
           response = "Command executed successfully.";
        }
        
        // Always show stdout if present
        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`;
        }
        
        // Always show EXIT_CODE if there was an error
        if (hasError) {
           response += `EXIT_CODE: ${exitCode !== null ? exitCode : "Unknown"}\n`;
        }

        return response.trim();

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

  const getWhitelistTool = tool({
    name: "get_whitelist",
    description: `Get the current allowed commands whitelist configuration.
    
Returns the comma-separated list of allowed commands as configured in the plugin settings. Only relevant when execution policy is set to "Allow Only".

TEST MODE: When the main run_shell_command tool has 'Test Mode' enabled, it will only check if commands are allowed and return status without executing them.`,
    parameters: {},
    implementation: async (_, { status }) => {
      const currentConfig = ctl.getPluginConfig(configSchematics);
      const policy = currentConfig.get("executionPolicy");
      const allowedStr = currentConfig.get("allowedCommands");
      
      status("Retrieving whitelist configuration...");
      
      if (policy !== "allow_only") {
        return `WHITELIST NOT ACTIVE: Current execution policy is '${policy}'. Whitelist only applies when policy is 'Allow Only'.`;
      }
      
      if (!allowedStr || allowedStr.trim() === "") {
        return "WHITELIST EMPTY: The Allowed Commands list is empty.";
      }
      
      const allowedList = allowedStr.split(",").map(s => s.trim()).filter(s => s.length > 0);
      return `WHITELIST (${allowedList.length} commands):\n${allowedList.join("\n")}`;
    },
  });

  const getBlacklistTool = tool({
    name: "get_blacklist",
    description: `Get the current forbidden commands blacklist configuration.
    
Returns the comma-separated list of blocked commands as configured in the plugin settings. This applies regardless of execution policy mode.

TEST MODE: When the main run_shell_command tool has 'Test Mode' enabled, it will only check if commands are allowed and return status without executing them.`,
    parameters: {},
    implementation: async (_, { status }) => {
      const currentConfig = ctl.getPluginConfig(configSchematics);
      const forbiddenStr = currentConfig.get("forbiddenCommands");
      
      status("Retrieving blacklist configuration...");
      
      if (!forbiddenStr || forbiddenStr.trim() === "") {
        return "BLACKLIST EMPTY: No forbidden commands configured.";
      }
      
      const forbiddenList = forbiddenStr.split(",").map(s => s.trim()).filter(s => s.length > 0);
      return `BLACKLIST (${forbiddenList.length} commands):\n${forbiddenList.join("\n")}`;
    },
  });

  return [executeCommandTool, getWhitelistTool, getBlacklistTool];
}