/pentestgpt/core/config.py

"""Configuration management for PentestGPT using Pydantic."""

from pathlib import Path
from typing import Any, Literal

from pydantic import Field
from pydantic_settings import BaseSettings, SettingsConfigDict


class PentestGPTConfig(BaseSettings):
    """Main configuration for PentestGPT."""

    model_config = SettingsConfigDict(
        env_file=".env",
        env_file_encoding="utf-8",
        case_sensitive=False,
        extra="ignore",
    )

    # LLM Configuration
    # Note: API key is optional - Claude Code manages its own configuration
    llm_model: str = Field(
        default="claude-sonnet-4-5-20250929", description="Claude model to use for the agent"
    )

    llm_api_key: str | None = Field(
        default=None, description="Optional API key (Claude Code manages its own config)"
    )

    llm_api_base: str | None = Field(default=None, description="Optional custom API base URL")

    # Agent Configuration
    max_iterations: int = Field(default=300, description="Maximum iterations for the agent")

    working_directory: Path = Field(
        default_factory=lambda: Path.cwd() / "workspace",
        description="Working directory for agent operations",
    )

    # Target Configuration
    target: str = Field(
        ...,  # Required
        description="Target for penetration testing (URL, IP, domain, or path)",
    )

    custom_instruction: str | None = Field(
        default=None, description="Optional custom instructions for the agent"
    )

    # Interface Configuration
    interface_mode: Literal["tui", "cli"] = Field(
        default="tui", description="Interface mode: TUI (interactive) or CLI (headless)"
    )

    verbose: bool = Field(default=True, description="Enable verbose output")

    # Permission Mode
    permission_mode: Literal["ask", "bypassPermissions"] = Field(
        default="bypassPermissions", description="Permission mode for Claude Code SDK"
    )

    def __init__(self, **data: Any) -> None:
        """Initialize configuration."""
        super().__init__(**data)

        # Create working directory if it doesn't exist
        # Ignore permission errors if directory already exists
        try:
            self.working_directory.mkdir(parents=True, exist_ok=True)
        except (PermissionError, OSError):
            # Directory already exists or we don't have permission to create it
            # This is fine if the directory is already available
            if not self.working_directory.exists():
                raise

    @property
    def system_prompt_path(self) -> Path:
        """Get path to system prompt file."""
        return Path(__file__).parent.parent / "prompts" / "pentesting.py"

    @classmethod
    def from_env(cls, **overrides: object) -> "PentestGPTConfig":
        """Create config from environment variables with optional overrides."""
        return cls(**overrides)


def load_config(**overrides: object) -> PentestGPTConfig:
    """
    Load configuration from environment with optional overrides.

    Args:
        **overrides: Keyword arguments to override config values

    Returns:
        PentestGPTConfig instance

    Example:
        >>> config = load_config(target="example.com", verbose=True)
    """
    # Create config with overrides
    # Note: API key is optional - Claude Code manages its own configuration
    return PentestGPTConfig.from_env(**overrides)