/pentestgpt/interface/main.py

#!/usr/bin/env python3
"""Main CLI entry point for PentestGPT."""

import argparse
import asyncio
import sys

from rich.console import Console
from rich.text import Text


def parse_arguments() -> argparse.Namespace:
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        prog="pentestgpt",
        description="PentestGPT - AI-Powered CTF Challenge Solver",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Interactive TUI mode (default) - HTB machine
  pentestgpt --target 10.10.11.234

  # Web challenge
  pentestgpt --target https://ctf.example.com/challenge1

  # With challenge context/hints
  pentestgpt --target 10.10.11.100 --instruction "Wordpress site, focus on plugin vulnerabilities"

  # Non-interactive mode
  pentestgpt --target challenge.htb --non-interactive

  # Custom model
  pentestgpt --target 10.10.11.50 --model claude-opus-4-20250514

For more information: https://github.com/yourusername/pentestgpt
        """,
    )

    parser.add_argument(
        "-t",
        "--target",
        type=str,
        required=True,
        help="Target CTF challenge or machine (URL, IP address, domain, or file path)",
    )

    parser.add_argument(
        "-i",
        "--instruction",
        type=str,
        help="Custom challenge context, hints, or instructions",
    )

    parser.add_argument(
        "-m",
        "--model",
        type=str,
        help="Claude model to use (default: claude-sonnet-4-5-20250929)",
    )

    parser.add_argument(
        "-n",
        "--non-interactive",
        action="store_true",
        help="Run in non-interactive mode (no TUI, exits on completion)",
    )

    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="Enable verbose output",
    )

    parser.add_argument(
        "-d",
        "--debug",
        action="store_true",
        help="Enable debug mode with detailed logging",
    )

    parser.add_argument(
        "--raw",
        action="store_true",
        help="Raw output mode for debugging (no TUI, no spinner, direct streaming)",
    )

    parser.add_argument(
        "--version",
        action="version",
        version="%(prog)s 1.0.0",
    )

    # Session management
    parser.add_argument(
        "-r",
        "--resume",
        action="store_true",
        help="Resume the most recent session for the target",
    )

    parser.add_argument(
        "--session-id",
        type=str,
        help="Resume a specific session by ID",
    )

    parser.add_argument(
        "--list-sessions",
        action="store_true",
        help="List available sessions and exit",
    )

    # Telemetry
    parser.add_argument(
        "--no-telemetry",
        action="store_true",
        help="Disable anonymous telemetry data collection",
    )

    return parser.parse_args()


def validate_environment() -> None:
    """
    Validate environment (optional checks).

    Note: API key is no longer required here as Claude Code manages
    its own configuration.
    """
    # No required validations - Claude Code handles API configuration
    pass


def print_banner() -> None:
    """Print welcome banner."""
    console = Console()

    banner = Text()
    banner.append("🚩 ", style="bold #6366f1")
    banner.append("PentestGPT", style="bold #6366f1")
    banner.append(" v1.0.0", style="dim")
    banner.append("\n")
    banner.append("AI-Powered CTF Challenge Solver", style="dim italic")

    console.print()
    console.print(banner)
    console.print()


async def run_cli_mode(args: argparse.Namespace) -> None:
    """Run in non-interactive CLI mode."""
    from pentestgpt.core.agent import run_pentest
    from pentestgpt.core.session import SessionStore

    console = Console()

    # Determine session to resume
    resume_session = args.session_id
    if args.resume and not resume_session:
        sessions = SessionStore()
        latest = sessions.get_latest(args.target)
        if latest:
            resume_session = latest.session_id
            console.print(f"[dim]Resuming session: {resume_session}[/]")

    console.print(f"[bold]Target:[/] {args.target}")
    if args.instruction:
        console.print(f"[bold]Challenge Context:[/] {args.instruction}")
    if args.debug:
        console.print("[dim]Debug mode: enabled[/]")
    console.print()

    # Allow up to 3 attempts
    max_attempts = 3
    attempt = 1
    custom_instruction = args.instruction
    total_cost = 0.0

    while attempt <= max_attempts:
        if attempt > 1:
            console.print(f"\n[bold cyan]Attempt {attempt}/{max_attempts}[/]")
            console.print("[dim]Adding persistence context...[/]\n")

        with console.status("[bold cyan]Solving CTF challenge...", spinner="dots"):
            result = await run_pentest(
                target=args.target,
                custom_instruction=custom_instruction,
                model=args.model,
                debug=args.debug,
                resume_session=resume_session if attempt == 1 else None,
            )

        total_cost += result.get("cost_usd", 0)

        console.print()
        if result["success"]:
            # Display flags found
            flags_found = result.get("flags_found", [])
            if flags_found:
                console.print(
                    f"[bold green]🚩 Challenge solved! {len(flags_found)} flag(s) found:[/]"
                )
                for flag_data in flags_found:
                    console.print(f"  • [bold cyan]{flag_data['flag']}[/]")

                console.print(f"\n{result['output']}")

                # Display session info
                session_id = result.get("session_id", "")
                if session_id:
                    console.print(f"\n[dim]Session: {session_id}[/]")

                if total_cost > 0:
                    console.print(f"[dim]Total cost: ${total_cost:.4f}[/]")

                # Success! Exit cleanly
                return

            else:
                # No flags found - this is a problem
                console.print("[bold red]⚠ CHALLENGE INCOMPLETE - NO FLAGS CAPTURED[/]")
                console.print("[yellow]The agent stopped without capturing flags.[/]")

                console.print(f"\n{result['output']}")

                if attempt < max_attempts:
                    console.print(f"\n[bold yellow]Attempts remaining: {max_attempts - attempt}[/]")
                    console.print(
                        "[dim]The agent will retry with stronger persistence instructions.[/]\n"
                    )

                    # Add feedback for next attempt
                    if custom_instruction:
                        custom_instruction += "\n\nPREVIOUS ATTEMPT FEEDBACK: You stopped without capturing flags. This is unacceptable. You MUST capture at least one flag. Try different approaches, enumerate harder, and do not stop until you have flags."
                    else:
                        custom_instruction = "IMPORTANT: The previous attempt failed to capture any flags. You MUST capture at least one flag. Do not stop until flags are found. Try all available techniques systematically."

                    attempt += 1
                    continue
                else:
                    console.print(f"\n[bold red]✗ Maximum attempts reached ({max_attempts})[/]")
                    console.print(
                        "[yellow]Consider providing more specific hints via --instruction[/]"
                    )

                    # Display session info for resumption
                    session_id = result.get("session_id", "")
                    if session_id:
                        console.print(
                            f"\n[dim]Session: {session_id} (resume with --session-id {session_id})[/]"
                        )

                    if total_cost > 0:
                        console.print(f"[dim]Total cost: ${total_cost:.4f}[/]")

                    sys.exit(1)
        else:
            console.print(
                f"[bold red]✗ Challenge failed:[/] {result.get('error', 'Unknown error')}"
            )
            sys.exit(1)


async def run_raw_mode(args: argparse.Namespace) -> None:
    """Run in raw CLI mode with streaming output for debugging."""
    from pentestgpt.core.config import load_config
    from pentestgpt.core.controller import AgentController
    from pentestgpt.core.events import Event, EventBus, EventType
    from pentestgpt.core.session import SessionStore

    # Print startup info
    print(f"[INFO] Target: {args.target}")
    if args.instruction:
        print(f"[INFO] Instruction: {args.instruction}")
    if args.model:
        print(f"[INFO] Model: {args.model}")
    if args.debug:
        print("[INFO] Debug mode: enabled")
    print("[INFO] Starting agent...", flush=True)

    # Set up event handlers that print directly to stdout
    events = EventBus.get()

    def on_message(event: Event) -> None:
        text = event.data.get("text", "")
        msg_type = event.data.get("type", "info")
        if text:
            print(f"[{msg_type.upper()}] {text}", flush=True)

    def on_tool(event: Event) -> None:
        status = event.data.get("status")
        name = event.data.get("name", "unknown")
        if status == "start":
            args_data = event.data.get("args", {})
            print(f"[TOOL] {name}: {args_data}", flush=True)
        else:
            print(f"[TOOL] {name} done", flush=True)

    def on_flag(event: Event) -> None:
        flag = event.data.get("flag", "")
        if flag:
            print(f"[FLAG] {flag}", flush=True)

    def on_state(event: Event) -> None:
        state = event.data.get("state", "")
        details = event.data.get("details", "")
        if details:
            print(f"[STATE] {state}: {details}", flush=True)
        else:
            print(f"[STATE] {state}", flush=True)

    events.subscribe(EventType.MESSAGE, on_message)
    events.subscribe(EventType.TOOL, on_tool)
    events.subscribe(EventType.FLAG_FOUND, on_flag)
    events.subscribe(EventType.STATE_CHANGED, on_state)

    # Determine session to resume
    resume_session = args.session_id
    if args.resume and not resume_session:
        sessions = SessionStore()
        latest = sessions.get_latest(args.target)
        if latest:
            resume_session = latest.session_id
            print(f"[INFO] Resuming session: {resume_session}", flush=True)

    # Build config
    config_kwargs = {"target": args.target}
    if args.instruction:
        config_kwargs["custom_instruction"] = args.instruction
    if args.model:
        config_kwargs["llm_model"] = args.model

    config = load_config(**config_kwargs)

    # Build task
    task = f"Solve this CTF challenge and capture the flag(s): {args.target}"
    if args.instruction:
        task += f"\n\nChallenge context: {args.instruction}"

    # Run agent with controller
    controller = AgentController(config)

    try:
        result = await controller.run(task, resume_session_id=resume_session)

        # Print final result
        if result.get("success"):
            flags = result.get("flags_found", [])
            cost = result.get("cost_usd", 0)
            session_id = result.get("session_id", "")
            print(
                f"[DONE] Flags: {len(flags)}, Cost: ${cost:.4f}, Session: {session_id}",
                flush=True,
            )
            if not flags:
                print("[WARN] No flags captured", flush=True)
                sys.exit(1)
        else:
            error = result.get("error", "Unknown error")
            print(f"[ERROR] {error}", flush=True)
            sys.exit(1)

    except Exception as e:
        print(f"[ERROR] {e!s}", flush=True)
        import traceback

        print(f"[TRACE] {traceback.format_exc()}", flush=True)
        sys.exit(1)


async def run_tui_mode(args: argparse.Namespace) -> None:
    """Run in interactive TUI mode."""
    from pentestgpt.core.session import SessionStore
    from pentestgpt.interface.tui import run_tui

    # Determine session to resume
    resume_session = args.session_id
    if args.resume and not resume_session:
        sessions = SessionStore()
        latest = sessions.get_latest(args.target)
        if latest:
            resume_session = latest.session_id

    await run_tui(
        target=args.target,
        custom_instruction=args.instruction,
        model=args.model,
        debug=args.debug,
        resume_session=resume_session,
    )


def list_sessions(target: str | None = None) -> None:
    """List available sessions."""
    from pentestgpt.core.session import SessionStore

    console = Console()
    sessions = SessionStore()
    session_list = sessions.list_sessions(target)

    if not session_list:
        console.print("[dim]No sessions found.[/]")
        return

    console.print(f"[bold]Sessions{f' for {target}' if target else ''}:[/]\n")
    console.print(f"{'ID':<10} {'Date':<18} {'Target':<25} {'Status':<12} {'Flags':<6}")
    console.print("-" * 75)

    for s in session_list:
        date_str = s.created_at.strftime("%Y-%m-%d %H:%M")
        target_str = s.target[:23] + ".." if len(s.target) > 25 else s.target
        flags_count = len(s.flags_found)
        console.print(
            f"{s.session_id:<10} {date_str:<18} {target_str:<25} {s.status.value:<12} {flags_count:<6}"
        )


def main() -> None:
    """Main entry point."""
    # Parse arguments
    args = parse_arguments()

    # Handle --list-sessions
    if args.list_sessions:
        list_sessions(args.target if hasattr(args, "target") else None)
        return

    # Validate environment
    validate_environment()

    # Initialize telemetry (enabled by default, use --no-telemetry to disable)
    from pentestgpt.core.langfuse import init_langfuse, shutdown_langfuse

    init_langfuse(disabled=args.no_telemetry)

    # Print banner in CLI mode
    if args.non_interactive:
        print_banner()

    # Run appropriate mode
    try:
        if args.raw:
            asyncio.run(run_raw_mode(args))
        elif args.non_interactive:
            asyncio.run(run_cli_mode(args))
        else:
            asyncio.run(run_tui_mode(args))

    except KeyboardInterrupt:
        console = Console()
        console.print("\n\n[yellow]Operation cancelled by user[/]")
        sys.exit(130)

    except Exception as e:
        console = Console()
        console.print(f"\n[bold red]Error:[/] {e!s}", style="red")
        # Always show traceback for debugging
        import traceback

        console.print("\n[dim]Traceback:[/]")
        console.print(traceback.format_exc())
        sys.exit(1)

    finally:
        shutdown_langfuse()


if __name__ == "__main__":
    main()