"""Main TUI application for PentestGPT using Textual framework.""" import asyncio import threading from pathlib import Path from typing import Any, ClassVar from rich.text import Text from textual import events from textual.app import App, ComposeResult from textual.binding import Binding from textual.containers import Grid, Vertical, VerticalScroll from textual.reactive import reactive from textual.screen import ModalScreen from textual.widgets import Button, Input, Label, Static from pentestgpt.core.controller import AgentController from pentestgpt.core.events import Event, EventBus, EventType from pentestgpt.interface.components.activity_feed import ActivityFeed from pentestgpt.interface.components.splash import SplashScreen class HelpScreen(ModalScreen[None]): """Modal screen showing keyboard shortcuts and help.""" def compose(self) -> ComposeResult: """Create the help dialog.""" yield Grid( Label("🚩 PentestGPT CTF Solver - Help", id="help_title"), Label( "F1 Help\n" "Ctrl+P Pause/Resume agent\n" "Ctrl+Q Quit\n" "Ctrl+C Quit\n" "↑/↓ Scroll feed\n" "PgUp/PgDn Fast scroll\n" "Enter Send instruction (when paused)", id="help_content", ), id="help_dialog", ) def on_key(self, _event: events.Key) -> None: """Close help on any key press.""" self.app.pop_screen() class QuitScreen(ModalScreen[None]): """Modal screen for quit confirmation.""" def compose(self) -> ComposeResult: """Create the quit confirmation dialog.""" yield Grid( Label("🚩 Quit PentestGPT CTF Solver?", id="quit_title"), Grid( Button("Yes", variant="error", id="btn_quit_confirm"), Button("No", variant="default", id="btn_quit_cancel"), id="quit_buttons", ), id="quit_dialog", ) def on_mount(self) -> None: """Focus cancel button by default.""" cancel_button = self.query_one("#btn_quit_cancel", Button) cancel_button.focus() def on_button_pressed(self, event: Button.Pressed) -> None: """Handle button clicks.""" if event.button.id == "btn_quit_confirm": self.app.exit() else: self.app.pop_screen() class PentestGPTApp(App[None]): """Main TUI application for PentestGPT CTF Solver.""" CSS_PATH = "styles.tcss" TITLE = "PentestGPT - CTF Challenge Solver" show_splash: reactive[bool] = reactive(default=True) agent_state: reactive[str] = reactive(default="idle") BINDINGS: ClassVar[list[Binding | tuple[str, str] | tuple[str, str, str]]] = [ Binding("f1", "toggle_help", "Help", priority=True), Binding("ctrl+p", "toggle_pause", "Pause/Resume", priority=True), Binding("ctrl+q", "request_quit", "Quit", priority=True), Binding("ctrl+c", "request_quit", "Quit", priority=True), ] def __init__( self, target: str, custom_instruction: str | None = None, model: str | None = None, debug: bool = False, resume_session: str | None = None, ) -> None: """Initialize the TUI app.""" super().__init__() self.target = target self.custom_instruction = custom_instruction self.model = model self._debug = debug self.resume_session = resume_session self._agent_thread: threading.Thread | None = None self._agent_stop_event = threading.Event() self._activity_feed: ActivityFeed | None = None self._controller: AgentController | None = None self._events: EventBus | None = None # Initialize in on_mount def _setup_event_handlers(self) -> None: """Subscribe to agent events for UI updates.""" if self._events is None: return self._events.subscribe(EventType.STATE_CHANGED, self._on_state_change) self._events.subscribe(EventType.MESSAGE, self._on_agent_message) self._events.subscribe(EventType.FLAG_FOUND, self._on_flag) self._events.subscribe(EventType.TOOL, self._on_tool) def _on_state_change(self, event: Event) -> None: """Handle agent state changes.""" state = event.data.get("state", "unknown") details = event.data.get("details", "") # Update reactive state (triggers watch_agent_state) self.call_from_thread(setattr, self, "agent_state", state) # Log state change if self._activity_feed and details: self.call_from_thread( self._activity_feed.add_message, f"Agent: {details}", "info", ) def _on_agent_message(self, event: Event) -> None: """Handle agent messages.""" if not self._activity_feed: return text = event.data.get("text", "") msg_type = event.data.get("type", "info") if text: self.call_from_thread( self._activity_feed.add_message, text, msg_type, ) def _on_flag(self, event: Event) -> None: """Handle flag found events.""" if not self._activity_feed: return flag = event.data.get("flag", "") if flag: self.call_from_thread( self._activity_feed.add_message, f"🚩 FLAG FOUND: {flag}", "success", ) def _on_tool(self, event: Event) -> None: """Handle tool events.""" if not self._activity_feed: return status = event.data.get("status", "") name = event.data.get("name", "") args = event.data.get("args", {}) result = event.data.get("result") if status == "start": self.call_from_thread( self._activity_feed.add_tool_execution, name, args, "running", None, ) elif status == "complete": self.call_from_thread( self._activity_feed.add_tool_execution, name, args, "completed", result, ) def watch_agent_state(self, state: str) -> None: """React to agent state changes - show/hide input field.""" try: user_input = self.query_one("#user_input", Input) # Show input when paused, hide otherwise user_input.display = state == "paused" if state == "paused": user_input.focus() except Exception: pass # Update status bar self._update_status_bar() def compose(self) -> ComposeResult: """Create the initial UI layout.""" if self.show_splash: yield SplashScreen(id="splash_screen") def watch_show_splash(self, show_splash: bool) -> None: """React to splash screen visibility changes.""" if not show_splash and hasattr(self, "_driver") and self._driver is not None: # Remove splash screen try: splash = self.query_one("#splash_screen") splash.remove() except Exception: pass # Build main interface self._build_main_interface() def _build_main_interface(self) -> None: """Build the main application interface.""" # Main container main_container = Vertical(id="main_container") self.mount(main_container) # Header header = self._create_header() main_container.mount(header) # Activity feed feed = ActivityFeed(id="activity_feed") self._activity_feed = feed content_area = VerticalScroll(feed, id="content_area") main_container.mount(content_area) # User input field (hidden by default, shown when paused) user_input = Input( placeholder="Type instruction and press Enter...", id="user_input", ) user_input.display = False main_container.mount(user_input) # Status bar status_bar = self._create_status_bar() main_container.mount(status_bar) def _create_header(self) -> Static: """Create the header bar.""" header_text = Text() header_text.append("🚩 ", style="bold #6366f1") header_text.append("PentestGPT", style="bold #6366f1") header_text.append(" CTF Solver", style="bold #6366f1") header_text.append(" v1.0", style="dim") header_text.append(" │ ", style="dim") header_text.append("Target: ", style="dim") header_text.append(self.target, style="bold") if self.model: header_text.append(" │ ", style="dim") header_text.append(f"Model: {self.model}", style="dim") header = Static(header_text, id="header") return header def _create_status_bar(self) -> Static: """Create the status bar.""" status_text = self._build_status_text() status_bar = Static(status_text, id="status_bar") return status_bar def _build_status_text(self) -> Text: """Build status bar text based on current state.""" status_text = Text() # Agent state status_text.append("Agent: ", style="dim") state_styles = { "idle": ("Idle", "dim"), "running": ("Running", "bold #6366f1"), "paused": ("PAUSED", "bold #f59e0b"), "completed": ("Completed", "bold #10b981"), "error": ("Error", "bold #ef4444"), } state_label, state_style = state_styles.get(self.agent_state, ("Unknown", "dim")) status_text.append(state_label, style=state_style) if self._debug: from pentestgpt.core.agent import DEBUG_LOG status_text.append(" │ ", style="dim") status_text.append("Debug: ", style="dim") status_text.append(str(DEBUG_LOG), style="bold #f59e0b") status_text.append(" │ ", style="dim") status_text.append("Ctrl+P", style="bold #8b5cf6") status_text.append(" Pause ", style="dim") status_text.append("F1", style="bold #8b5cf6") status_text.append(" Help ", style="dim") status_text.append("Ctrl+Q", style="bold #8b5cf6") status_text.append(" Quit", style="dim") return status_text def _update_status_bar(self) -> None: """Update status bar with current state.""" try: status_bar = self.query_one("#status_bar", Static) status_bar.update(self._build_status_text()) except Exception: pass def on_mount(self) -> None: """Handle app mount - show splash then start agent.""" # Initialize event bus after app is mounted self._events = EventBus.get() self._setup_event_handlers() # Show splash for 2.5 seconds self.set_timer(2.5, self._hide_splash) def _hide_splash(self) -> None: """Hide splash screen and start main app.""" self.show_splash = False # Start agent after splash self.call_later(self._start_agent) def _start_agent(self) -> None: """Start the agent in a background thread using AgentController.""" from pentestgpt.core.config import load_config if self._activity_feed: self._activity_feed.add_message( "🚩 CTF Solver initializing...", message_type="info", ) def run_agent() -> None: """Run the agent with controller.""" try: # Build config config_kwargs: dict[str, Any] = { "target": self.target, "working_directory": Path("/workspace"), } if self.custom_instruction: config_kwargs["custom_instruction"] = self.custom_instruction if self.model: config_kwargs["llm_model"] = self.model config = load_config(**config_kwargs) # Create controller self._controller = AgentController(config) # Build task task = f"Solve this CTF challenge and capture the flag(s): {self.target}" if self.custom_instruction: task += f"\n\nChallenge context: {self.custom_instruction}" # Run agent with controller loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) result = loop.run_until_complete( self._controller.run(task, resume_session_id=self.resume_session) ) loop.close() # Show final result summary if self._activity_feed: if result.get("success"): flags = result.get("flags_found", []) cost = result.get("cost_usd", 0) session_id = result.get("session_id", "") if flags: self.call_from_thread( self._activity_feed.add_message, f"🚩 Challenge complete! {len(flags)} flag(s) | Cost: ${cost:.4f} | Session: {session_id}", "success", ) else: self.call_from_thread( self._activity_feed.add_message, f"⚠ Challenge ended without flags | Cost: ${cost:.4f} | Session: {session_id}", "warning", ) else: error = result.get("error", "Unknown error") self.call_from_thread( self._activity_feed.add_message, f"✗ Agent error: {error}", "error", ) except Exception as e: if self._activity_feed: self.call_from_thread( self._activity_feed.add_message, f"✗ Agent error: {e!s}", "error", ) self._agent_thread = threading.Thread(target=run_agent, daemon=True) self._agent_thread.start() def on_input_submitted(self, event: Input.Submitted) -> None: """Handle user input submission.""" if event.input.id == "user_input": text = event.value.strip() if text and self._events: # Emit user input event self._events.emit_input(text) if self._activity_feed: self._activity_feed.add_message( f"📝 Instruction queued: {text[:50]}...", message_type="info", ) event.input.value = "" def action_toggle_pause(self) -> None: """Toggle pause/resume state (Ctrl+P).""" if self.show_splash or not self._controller: return if self.agent_state == "paused": self._controller.resume() elif self.agent_state == "running": self._controller.pause() def action_toggle_help(self) -> None: """Show/hide help screen.""" if self.show_splash: return if isinstance(self.screen, HelpScreen): self.pop_screen() else: self.push_screen(HelpScreen()) def action_request_quit(self) -> None: """Show quit confirmation.""" if self.show_splash: self.exit() return self.push_screen(QuitScreen()) def on_unmount(self) -> None: """Cleanup when app closes.""" # Signal agent to stop self._agent_stop_event.set() # Wait for agent thread if self._agent_thread and self._agent_thread.is_alive(): self._agent_thread.join(timeout=2.0) async def run_tui( target: str, custom_instruction: str | None = None, model: str | None = None, debug: bool = False, resume_session: str | None = None, ) -> None: """Run the TUI application. Args: target: Target URL/IP/domain custom_instruction: Optional challenge context model: Optional model override debug: Enable debug mode resume_session: Optional session ID to resume """ app = PentestGPTApp( target=target, custom_instruction=custom_instruction, model=model, debug=debug, resume_session=resume_session, ) await app.run_async() # For testing if __name__ == "__main__": import asyncio asyncio.run(run_tui("example.com"))