src/tools/shell.py

"""
Shell tool — shell command execution.

Allows agents to execute commands in the system shell.
Supports timeouts and output size limits.
"""

import subprocess
import sys
from typing import Any

from .base import BaseTool, ToolResult


class ShellTool(BaseTool):
    """
    Tool for executing shell commands.

    Security:
        - Commands are executed in a subprocess with a timeout
        - Output is size-limited to prevent overflow
        - shell=True is not used on Windows for security

    Example:
        tool = ShellTool(timeout=30, max_output_size=4096)
        result = tool.execute(command="ls -la")

        if result.success:
            print(result.output)
        else:
            print(f"Error: {result.error}")

    """

    def __init__(
        self,
        timeout: int = 30,
        max_output_size: int = 8192,
        working_dir: str | None = None,
        allowed_commands: list[str] | None = None,
    ):
        """
        Create ShellTool.

        Args:
            timeout: Maximum command execution time in seconds.
            max_output_size: Maximum output size in bytes.
            working_dir: Working directory for commands.
            allowed_commands: Whitelist of allowed commands (None = all).

        """
        self._timeout = timeout
        self._max_output_size = max_output_size
        self._working_dir = working_dir
        self._allowed_commands = set(allowed_commands) if allowed_commands else None

    @property
    def name(self) -> str:
        return "shell"

    @property
    def description(self) -> str:
        return (
            "Execute a shell command and return its output. "
            "Use for system operations, file manipulation, or running scripts."
        )

    @property
    def parameters_schema(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "command": {
                    "type": "string",
                    "description": "The shell command to execute",
                },
            },
            "required": ["command"],
        }

    def _is_command_allowed(self, command: str) -> bool:
        """Check whether the command is allowed."""
        if self._allowed_commands is None:
            return True

        # Extract the first word (command name)
        cmd_name = command.strip().split()[0] if command.strip() else ""
        return cmd_name in self._allowed_commands

    def execute(self, command: str = "", **_kwargs: Any) -> ToolResult:
        """
        Execute a shell command.

        Args:
            command: Command to execute.

        Returns:
            ToolResult with the command output or error.

        """
        if not command:
            return ToolResult(
                tool_name=self.name,
                success=False,
                error="No command provided",
            )

        if not self._is_command_allowed(command):
            return ToolResult(
                tool_name=self.name,
                success=False,
                error=f"Command not allowed: {command.split()[0]}",
            )

        try:
            # Determine shell based on OS
            if sys.platform == "win32":
                # On Windows use cmd.exe
                result = subprocess.run(
                    command,
                    shell=True,
                    capture_output=True,
                    text=True,
                    timeout=self._timeout,
                    cwd=self._working_dir,
                    check=False,
                )
            else:
                # On Unix use /bin/sh
                result = subprocess.run(
                    command,
                    shell=True,
                    executable="/bin/sh",
                    capture_output=True,
                    text=True,
                    timeout=self._timeout,
                    cwd=self._working_dir,
                    check=False,
                )

            # Merge stdout and stderr
            output = result.stdout
            if result.stderr:
                output += f"\n[stderr]\n{result.stderr}"

            # Limit output size
            if len(output) > self._max_output_size:
                output = output[: self._max_output_size] + "\n... (output truncated)"

            if result.returncode != 0:
                return ToolResult(
                    tool_name=self.name,
                    success=False,
                    output=output,
                    error=f"Command exited with code {result.returncode}",
                )

            return ToolResult(
                tool_name=self.name,
                success=True,
                output=output.strip() if output else "(no output)",
            )

        except subprocess.TimeoutExpired:
            return ToolResult(
                tool_name=self.name,
                success=False,
                error=f"Command timed out after {self._timeout} seconds",
            )
        except FileNotFoundError:
            return ToolResult(
                tool_name=self.name,
                success=False,
                error="Command not found",
            )
        except (OSError, ValueError) as e:
            return ToolResult(
                tool_name=self.name,
                success=False,
                error=f"Execution error: {e}",
            )