pentestagent/ui/menu_system.py

"""Menu system and user interface components for GHOSTCREW."""

from typing import Optional, List, Tuple
from colorama import Fore, Style
from config.constants import (
    MAIN_MENU_TITLE, INTERACTIVE_OPTION, AUTOMATED_OPTION, 
    EXIT_OPTION, MULTI_LINE_PROMPT, MULTI_LINE_END_MARKER
)


class MenuSystem:
    """Handles all menu displays and user input for GHOSTCREW."""
    
    @staticmethod
    def display_main_menu(workflows_available: bool, has_connected_servers: bool) -> None:
        """Display the main application menu."""
        print(f"\n{MAIN_MENU_TITLE}")
        print(f"1. {INTERACTIVE_OPTION}")
        
        # Check if automated mode should be available
        if workflows_available and has_connected_servers:
            print(f"2. {AUTOMATED_OPTION}")
        elif workflows_available and not has_connected_servers:
            print(f"2. {Fore.LIGHTBLACK_EX}Automated Penetration Testing (requires MCP tools){Style.RESET_ALL}")
        else:
            print(f"2. {Fore.LIGHTBLACK_EX}Automated Penetration Testing (workflows.py not found){Style.RESET_ALL}")
        
        print(f"3. {EXIT_OPTION}")
    
    @staticmethod
    def get_menu_choice(max_option: int = 3) -> str:
        """Get user's menu selection."""
        return input(f"\n{Fore.GREEN}Select mode (1-{max_option}): {Style.RESET_ALL}").strip()
    
    @staticmethod
    def display_interactive_mode_intro() -> None:
        """Display introduction for interactive chat mode."""
        print(f"\n{Fore.CYAN}INTERACTIVE CHAT MODE{Style.RESET_ALL}")
        print(f"{Fore.WHITE}Type your questions or commands. Use 'multi' for multi-line input.{Style.RESET_ALL}")
        print(f"{Fore.WHITE}Type 'menu' to return to main menu.{Style.RESET_ALL}\n")
    
    @staticmethod
    def get_user_input() -> str:
        """Get user input with prompt."""
        print(f"\n{Fore.GREEN}[>]{Style.RESET_ALL} ", end="")
        return input().strip()
    
    @staticmethod
    def get_multi_line_input() -> Optional[str]:
        """Get multi-line input from user."""
        print(f"{Fore.CYAN}Entering multi-line mode. Type your query across multiple lines.{Style.RESET_ALL}")
        print(f"{Fore.CYAN}Press Enter on empty line to submit.{Style.RESET_ALL}")
        
        lines = []
        while True:
            line = input()
            if line == "":  # Empty line ends input
                break
            lines.append(line)
        
        # Only proceed if they actually entered something
        if not lines:
            print(f"{Fore.YELLOW}No query entered in multi-line mode.{Style.RESET_ALL}")
            return None
        
        return "\n".join(lines)
    
    @staticmethod
    def display_no_query_message() -> None:
        """Display message when no query is entered."""
        print(f"{Fore.YELLOW}No query entered. Please type your question.{Style.RESET_ALL}")
    
    @staticmethod
    def display_ready_message() -> None:
        """Display ready for next query message."""
        print(f"\n{Fore.CYAN}Ready for next query. Type 'quit', 'multi' for multi-line, or 'menu' for main menu.{Style.RESET_ALL}")
    
    @staticmethod
    def display_exit_message() -> None:
        """Display exit message."""
        print(f"\n{Fore.CYAN}Thank you for using GHOSTCREW, exiting...{Style.RESET_ALL}")
    
    @staticmethod
    def display_workflow_requirements_message() -> None:
        """Display message about automated workflow requirements."""
        print(f"\n{Fore.YELLOW}Automated penetration testing requires MCP tools to be configured and connected.{Style.RESET_ALL}")
        print(f"{Fore.WHITE}Without real security tools, the AI would only generate simulated responses.{Style.RESET_ALL}")
        print(f"{Fore.WHITE}Please configure MCP tools to use this feature.{Style.RESET_ALL}")
        input(f"{Fore.CYAN}Press Enter to continue...{Style.RESET_ALL}")
    
    @staticmethod
    def get_workflow_target() -> Optional[str]:
        """Get target input for workflow execution."""
        target = input(f"{Fore.YELLOW}Enter target (IP, domain, or network): {Style.RESET_ALL}").strip()
        if not target:
            print(f"{Fore.RED}Target is required.{Style.RESET_ALL}")
            return None
        return target
    
    @staticmethod
    def confirm_workflow_execution(workflow_name: str, target: str) -> bool:
        """Confirm workflow execution with user."""
        confirm = input(f"{Fore.YELLOW}Execute '{workflow_name}' on {target}? (yes/no): {Style.RESET_ALL}").strip().lower()
        return confirm == 'yes'
    
    @staticmethod
    def display_workflow_cancelled() -> None:
        """Display workflow cancelled message."""
        print(f"{Fore.YELLOW}Workflow cancelled.{Style.RESET_ALL}")
    
    @staticmethod
    def display_workflow_completed() -> None:
        """Display workflow completion message."""
        print(f"\n{Fore.GREEN}Workflow completed successfully!{Style.RESET_ALL}")
    
    @staticmethod
    def ask_generate_report() -> bool:
        """Ask if user wants to generate a report."""
        response = input(f"\n{Fore.CYAN}Generate markdown report? (yes/no): {Style.RESET_ALL}").strip().lower()
        return response == 'yes'
    
    @staticmethod
    def ask_save_raw_history() -> bool:
        """Ask if user wants to save raw conversation history."""
        response = input(f"{Fore.YELLOW}Save raw conversation history? (yes/no, default: no): {Style.RESET_ALL}").strip().lower()
        return response == 'yes'
    
    @staticmethod
    def display_report_generated(report_path: str) -> None:
        """Display report generation success message."""
        print(f"\n{Fore.GREEN}Report generated: {report_path}{Style.RESET_ALL}")
        print(f"{Fore.CYAN}Open the markdown file in any markdown viewer for best formatting{Style.RESET_ALL}")
    
    @staticmethod
    def display_report_error(error: Exception) -> None:
        """Display report generation error message."""
        print(f"\n{Fore.RED}Error generating report: {error}{Style.RESET_ALL}")
        print(f"{Fore.YELLOW}Raw workflow data is still available in conversation history{Style.RESET_ALL}")
    
    @staticmethod
    def display_invalid_choice() -> None:
        """Display invalid choice message."""
        print(f"{Fore.RED}Invalid choice. Please select a valid option.{Style.RESET_ALL}")
    
    @staticmethod
    def display_invalid_input() -> None:
        """Display invalid input message."""
        print(f"{Fore.RED}Invalid input. Please enter a number.{Style.RESET_ALL}")
    
    @staticmethod
    def display_operation_cancelled() -> None:
        """Display operation cancelled message."""
        print(f"\n{Fore.YELLOW}Operation cancelled.{Style.RESET_ALL}")
    
    @staticmethod
    def press_enter_to_continue() -> None:
        """Wait for user to press enter."""
        input(f"\n{Fore.CYAN}Press Enter to continue...{Style.RESET_ALL}")