/pentestgpt/tools/base.py

"""Base classes for extensible tool framework."""

from abc import ABC, abstractmethod
from typing import Any


class BaseTool(ABC):
    """Base class for all tools in PentestGPT."""

    def __init__(self, name: str, description: str) -> None:
        """
        Initialize a tool.

        Args:
            name: Tool name (e.g., "terminal_execute")
            description: Human-readable description
        """
        self.name = name
        self.description = description

    @abstractmethod
    async def execute(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
        """
        Execute the tool with given arguments.

        Returns:
            Dictionary with execution results including:
            - success: bool
            - result: Any
            - error: Optional[str]
        """
        pass

    def to_dict(self) -> dict[str, Any]:
        """Convert tool execution info to dictionary for tracer."""
        return {
            "tool_name": self.name,
            "description": self.description,
        }


class TerminalTool(BaseTool):
    """Tool for executing terminal commands (via Claude Code)."""

    def __init__(self) -> None:
        """Initialize terminal tool."""
        super().__init__(
            name="terminal_execute", description="Execute shell commands for penetration testing"
        )

    async def execute(self, command: str = "", **kwargs: Any) -> dict[str, Any]:
        """
        Execute a terminal command.

        Note: For MVP, this is handled by Claude Code's built-in Bash tool.
        This class serves as a placeholder for future enhancements.

        Args:
            command: Shell command to execute
            **kwargs: Additional arguments

        Returns:
            Execution results
        """
        # In the MVP, Claude Code handles command execution
        # This is a placeholder for future direct execution capability
        return {
            "success": True,
            "command": command,
            "result": "Executed via Claude Code",
        }