hackagent 0.3.1__py3-none-any.whl

hackagent/cli/tui/widgets/logs.py

@@ -0,0 +1,435 @@
+# Copyright 2025 - AI4I. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Attack Log Viewer Component
+
+A reusable Textual widget for displaying live attack execution logs
+with syntax highlighting, auto-scrolling, and filtering capabilities.
+"""
+
+from textual.app import ComposeResult
+from textual.containers import Container
+from textual.widgets import Button, RichLog, Static
+
+
+class AttackLogViewer(Container):
+    """
+    A container widget for displaying attack execution logs in real-time.
+
+    This component provides:
+    - Live log streaming with syntax highlighting
+    - Color-coded log levels (INFO, WARNING, ERROR)
+    - Auto-scroll to latest logs
+    - Manual scroll capability
+    - Clear logs functionality
+    - Export logs to file
+    """
+
+    DEFAULT_CSS = """
+    AttackLogViewer {
+        border: solid $primary;
+        padding: 0;
+    }
+
+    AttackLogViewer .log-header {
+        dock: top;
+        height: 3;
+        background: $panel;
+        padding: 0 1;
+        content-align: center middle;
+    }
+
+    AttackLogViewer .log-controls {
+        dock: top;
+        height: 3;
+        background: $surface;
+        padding: 0 1;
+        layout: horizontal;
+    }
+
+    AttackLogViewer RichLog {
+        background: $surface;
+        border: none;
+        padding: 1;
+        height: 1fr;
+        width: 100%;
+    }
+
+    AttackLogViewer Button {
+        margin: 0 1;
+    }
+    """
+
+    def __init__(
+        self,
+        title: str = "Attack Execution Logs",
+        show_controls: bool = True,
+        max_lines: int = 1000,
+        **kwargs,
+    ):
+        """
+        Initialize the log viewer.
+
+        Args:
+            title: Title to display in the header
+            show_controls: Whether to show control buttons
+            max_lines: Maximum number of log lines to retain
+            **kwargs: Additional keyword arguments for Container
+        """
+        super().__init__(**kwargs)
+        self.log_title = title
+        self.show_controls = show_controls
+        self.max_lines = max_lines
+        self._auto_scroll = True
+        self._line_count = 0  # Track line count internally
+        self._log_buffer: list[str] = []  # Store log messages for copying
+        self._log_file = None  # File handle for continuous logging
+        self._log_file_path = None  # Path to the log file
+        self._initialize_log_file()
+
+    def compose(self) -> ComposeResult:
+        """Compose the log viewer layout."""
+        # Header
+        yield Static(
+            f"[bold cyan]{self.log_title}[/bold cyan]",
+            classes="log-header",
+        )
+
+        # Control buttons (optional)
+        if self.show_controls:
+            with Container(classes="log-controls"):
+                yield Button("Clear Logs", id="clear-logs", variant="default")
+                yield Button("Auto-scroll: ON", id="toggle-scroll", variant="primary")
+                yield Static("", id="log-count")
+                yield Static("", id="log-file-path", classes="log-file-info")
+
+        # Log display area
+        rich_log = RichLog(
+            highlight=True,
+            markup=True,
+            max_lines=self.max_lines,
+            wrap=True,
+            id="attack-log-display",
+        )
+        yield rich_log
+
+    def _initialize_log_file(self) -> None:
+        """Initialize log file for continuous writing."""
+        try:
+            import os
+            from datetime import datetime
+
+            # Create logs directory
+            log_dir = os.path.join(os.getcwd(), "logs")
+            os.makedirs(log_dir, exist_ok=True)
+
+            # Create log file with timestamp
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            self._log_file_path = os.path.join(log_dir, f"attack_logs_{timestamp}.log")
+            self._log_file = open(
+                self._log_file_path, "w", buffering=1
+            )  # Line buffered
+
+        except Exception:
+            # If file creation fails, just continue without file logging
+            self._log_file = None
+            self._log_file_path = None
+
+    def on_mount(self) -> None:
+        """Called when the widget is mounted."""
+        self.update_log_count(0)
+        if self._log_file_path and self.show_controls:
+            self._update_log_file_path_display()
+
+    def on_button_pressed(self, event: Button.Pressed) -> None:
+        """Handle button press events."""
+        if event.button.id == "clear-logs":
+            self.clear_logs()
+        elif event.button.id == "toggle-scroll":
+            self.toggle_auto_scroll()
+
+    def add_log(self, message: str, level: str = "INFO") -> None:
+        """
+        Add a log message to the viewer with appropriate styling.
+
+        Args:
+            message: The log message to display
+            level: Log level (INFO, WARNING, ERROR, DEBUG)
+        """
+        log_widget = self.query_one("#attack-log-display", RichLog)
+
+        # Color code based on log level
+        level_colors = {
+            "DEBUG": "dim",
+            "INFO": "cyan",
+            "WARNING": "yellow",
+            "ERROR": "bold red",
+            "CRITICAL": "bold red on white",
+        }
+
+        color = level_colors.get(level, "white")
+
+        # Format the message with color
+        if level in ["ERROR", "CRITICAL"]:
+            formatted_message = f"[{color}]🔴 {message}[/{color}]"
+        elif level == "WARNING":
+            formatted_message = f"[{color}]⚠️  {message}[/{color}]"
+        elif level == "DEBUG":
+            formatted_message = f"[{color}]🔍 {message}[/{color}]"
+        else:  # INFO and default
+            formatted_message = f"[{color}]{message}[/{color}]"
+
+        # Add to log display
+        log_widget.write(formatted_message)
+
+        # Store in buffer for copying (strip Rich markup)
+        plain_message = message  # Store the original message without formatting
+        log_entry = f"[{level}] {plain_message}"
+        self._log_buffer.append(log_entry)
+
+        # Write to log file immediately
+        if self._log_file:
+            try:
+                self._log_file.write(log_entry + "\n")
+                self._log_file.flush()  # Ensure it's written immediately
+            except Exception:
+                pass  # Silently continue if file write fails
+
+        # Auto-scroll to bottom if enabled
+        if self._auto_scroll:
+            log_widget.scroll_end(animate=False)
+
+        # Update log count
+        self._line_count += 1
+        self.update_log_count(self._line_count)
+
+    def add_step_header(self, step_name: str, step_number: int = 0) -> None:
+        """
+        Add a prominent step header to visually separate pipeline steps.
+
+        Args:
+            step_name: Name of the step
+            step_number: Step number (0 for no number)
+        """
+        log_widget = self.query_one("#attack-log-display", RichLog)
+
+        # Create a visual separator
+        separator = "─" * 60
+        if step_number > 0:
+            header = f"\n[bold magenta]{separator}\n🎯 STEP {step_number}: {step_name}\n{separator}[/bold magenta]\n"
+        else:
+            header = f"\n[bold magenta]{separator}\n🎯 {step_name}\n{separator}[/bold magenta]\n"
+
+        log_widget.write(header)
+
+        if self._auto_scroll:
+            log_widget.scroll_end(animate=False)
+
+    def clear_logs(self) -> None:
+        """Clear all log messages from the viewer."""
+        log_widget = self.query_one("#attack-log-display", RichLog)
+        log_widget.clear()
+        self._line_count = 0
+        self._log_buffer.clear()
+
+        # Close current log file and create a new one
+        if self._log_file:
+            try:
+                self._log_file.close()
+            except Exception:
+                pass
+
+        self._initialize_log_file()
+        if self._log_file_path:
+            self._update_log_file_path_display()
+
+        self.update_log_count(0)
+
+    def copy_logs(self) -> None:
+        """Copy all log messages to clipboard or save to file."""
+        if not self._log_buffer:
+            return
+
+        log_text = "\n".join(self._log_buffer)
+
+        # Try multiple clipboard methods
+        copied = False
+
+        # Method 1: Try subprocess clipboard tools first (more reliable in containers/SSH)
+        try:
+            import subprocess
+            import platform
+
+            system = platform.system()
+            if system == "Linux":
+                # Try xclip first, then xsel
+                try:
+                    subprocess.run(
+                        ["xclip", "-selection", "clipboard"],
+                        input=log_text.encode(),
+                        check=True,
+                        stderr=subprocess.DEVNULL,
+                        timeout=2,
+                    )
+                    copied = True
+                except (
+                    FileNotFoundError,
+                    subprocess.CalledProcessError,
+                    subprocess.TimeoutExpired,
+                ):
+                    try:
+                        subprocess.run(
+                            ["xsel", "--clipboard", "--input"],
+                            input=log_text.encode(),
+                            check=True,
+                            stderr=subprocess.DEVNULL,
+                            timeout=2,
+                        )
+                        copied = True
+                    except (
+                        FileNotFoundError,
+                        subprocess.CalledProcessError,
+                        subprocess.TimeoutExpired,
+                    ):
+                        pass
+            elif system == "Darwin":  # macOS
+                subprocess.run(
+                    ["pbcopy"], input=log_text.encode(), check=True, timeout=2
+                )
+                copied = True
+            elif system == "Windows":
+                subprocess.run(["clip"], input=log_text.encode(), check=True, timeout=2)
+                copied = True
+
+            if copied:
+                pass
+        except Exception:
+            pass
+
+        # Method 2: Try pyperclip as fallback (if subprocess failed)
+        if not copied:
+            try:
+                import pyperclip
+
+                pyperclip.copy(log_text)
+                copied = True
+            except ImportError:
+                pass
+            except Exception:
+                pass
+
+        # Fallback: Save to file
+        if not copied:
+            try:
+                import os
+
+                # Create a more permanent location in the project
+                log_dir = os.path.join(os.getcwd(), "logs")
+                os.makedirs(log_dir, exist_ok=True)
+
+                from datetime import datetime
+
+                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+                log_file = os.path.join(log_dir, f"attack_logs_{timestamp}.log")
+
+                with open(log_file, "w") as f:
+                    f.write(log_text)
+            except Exception:
+                pass
+
+    def view_in_pager(self) -> None:
+        """View logs in a pager (less) for easy selection and navigation."""
+        if not self._log_buffer:
+            return
+
+        try:
+            import tempfile
+            import subprocess
+            import os
+
+            # Save to temporary file
+            log_text = "\n".join(self._log_buffer)
+            temp_file = tempfile.NamedTemporaryFile(
+                mode="w", suffix=".log", delete=False
+            )
+            temp_file.write(log_text)
+            temp_file.close()
+
+            # Suspend the TUI and open in pager
+            self.app.suspend()
+
+            # Try less first (with mouse support), fall back to more
+            pager = os.environ.get("PAGER", "less")
+            if pager == "less":
+                # Enable mouse, color, and exit if content fits on screen
+                subprocess.run(["less", "-R", "-X", "--mouse", temp_file.name])
+            else:
+                subprocess.run([pager, temp_file.name])
+
+            # Clean up
+            os.unlink(temp_file.name)
+
+            # Resume the TUI
+            self.app.refresh()
+
+        except Exception:
+            self.app.refresh()  # Make sure we resume even on error
+            pass
+
+    def toggle_auto_scroll(self) -> None:
+        """Toggle automatic scrolling to latest logs."""
+        self._auto_scroll = not self._auto_scroll
+        button = self.query_one("#toggle-scroll", Button)
+        button.label = f"Auto-scroll: {'ON' if self._auto_scroll else 'OFF'}"
+        button.variant = "primary" if self._auto_scroll else "default"
+
+    def update_log_count(self, count: int) -> None:
+        """
+        Update the log count display.
+
+        Args:
+            count: Number of log lines currently displayed
+        """
+        if self.show_controls:
+            count_widget = self.query_one("#log-count", Static)
+            count_widget.update(f"[dim]Lines: {count}/{self.max_lines}[/dim]")
+
+    def _update_log_file_path_display(self) -> None:
+        """Update the log file path display."""
+        if self.show_controls and self._log_file_path:
+            try:
+                path_widget = self.query_one("#log-file-path", Static)
+                path_widget.update(f"[dim]📄 Logs: {self._log_file_path}[/dim]")
+            except Exception:
+                pass  # Widget might not be mounted yet
+
+    def get_log_text(self) -> str:
+        """
+        Get all log text as a plain string (for export).
+
+        Returns:
+            All log messages as plain text
+        """
+        return "\n".join(self._log_buffer)
+
+    def load_logs_from_buffer(self, buffer: list[tuple[str, str]]) -> None:
+        """
+        Load logs from a buffer (e.g., from TUILogHandler).
+
+        Args:
+            buffer: List of (message, level) tuples
+        """
+        for message, level in buffer:
+            self.add_log(message, level)  # add_log will handle line count

hackagent/cli/utils.py

@@ -0,0 +1,276 @@
+# Copyright 2025 - AI4I. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+CLI Utilities
+
+Common utilities for the HackAgent CLI including error handling,
+formatting, and helper functions.
+"""
+
+import functools
+import json
+from pathlib import Path
+from typing import Any, Dict
+
+import click
+from rich.console import Console
+from rich.panel import Panel
+from rich.table import Table
+from rich.text import Text
+from rich.traceback import Traceback
+
+from hackagent.errors import ApiError, HackAgentError
+
+console = Console()
+
+
+def handle_errors(func):
+    """Decorator for consistent error handling across CLI commands"""
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except HackAgentError as e:
+            console.print(f"[bold red]❌ HackAgent Error: {str(e)}")
+            if console._environ.get("HACKAGENT_DEBUG"):
+                console.print(Traceback())
+            raise click.ClickException(str(e))
+        except ApiError as e:
+            console.print(f"[bold red]❌ API Error: {str(e)}")
+            if console._environ.get("HACKAGENT_DEBUG"):
+                console.print(Traceback())
+            raise click.ClickException(str(e))
+        except ValueError as e:
+            console.print(f"[bold red]❌ Configuration Error: {str(e)}")
+            raise click.ClickException(str(e))
+        except FileNotFoundError as e:
+            console.print(f"[bold red]❌ File Not Found: {str(e)}")
+            raise click.ClickException(str(e))
+        except Exception as e:
+            console.print(f"[bold red]❌ Unexpected error: {str(e)}")
+            if console._environ.get("HACKAGENT_DEBUG"):
+                console.print(Traceback())
+            raise click.ClickException(str(e))
+
+    return wrapper
+
+
+def load_config_file(path: str) -> Dict[str, Any]:
+    """Load configuration from YAML or JSON file"""
+    config_path = Path(path)
+
+    if not config_path.exists():
+        raise FileNotFoundError(f"Configuration file not found: {path}")
+
+    try:
+        with open(config_path) as f:
+            if config_path.suffix.lower() in [".yaml", ".yml"]:
+                try:
+                    import yaml
+
+                    return yaml.safe_load(f) or {}
+                except ImportError:
+                    raise click.ClickException(
+                        "PyYAML required for YAML config files. Install with: pip install pyyaml"
+                    )
+            else:
+                return json.load(f)
+    except json.JSONDecodeError as e:
+        raise click.ClickException(f"Invalid JSON in config file {path}: {e}")
+    except Exception as e:
+        raise click.ClickException(f"Failed to load config file {path}: {e}")
+
+
+def display_results_table(results: Any, title: str = "Results") -> None:
+    """Display results in a formatted table"""
+    import pandas as pd
+
+    if isinstance(results, pd.DataFrame):
+        if results.empty:
+            console.print(f"[yellow]ℹ️ No {title.lower()} found")
+            return
+
+        table = Table(title=title, show_header=True, header_style="bold cyan")
+
+        # Add columns
+        for column in results.columns:
+            table.add_column(str(column))
+
+        # Add rows (limit to first 20 for display)
+        display_results = results.head(20)
+        for _, row in display_results.iterrows():
+            table.add_row(*[str(value) for value in row])
+
+        console.print(table)
+
+        if len(results) > 20:
+            console.print(f"[dim]... and {len(results) - 20} more rows")
+
+    elif isinstance(results, list):
+        if not results:
+            console.print(f"[yellow]ℹ️ No {title.lower()} found")
+            return
+
+        # Try to create table from list of dicts
+        if results and isinstance(results[0], dict):
+            table = Table(title=title, show_header=True, header_style="bold cyan")
+
+            # Get all unique keys for columns
+            all_keys = set()
+            for item in results:
+                all_keys.update(item.keys())
+
+            # Add columns
+            for key in sorted(all_keys):
+                table.add_column(str(key))
+
+            # Add rows (limit to first 20)
+            for item in results[:20]:
+                row_values = []
+                for key in sorted(all_keys):
+                    value = item.get(key, "")
+                    row_values.append(str(value))
+                table.add_row(*row_values)
+
+            console.print(table)
+
+            if len(results) > 20:
+                console.print(f"[dim]... and {len(results) - 20} more rows")
+        else:
+            # Simple list display
+            for i, item in enumerate(results[:20], 1):
+                console.print(f"{i}. {item}")
+
+            if len(results) > 20:
+                console.print(f"[dim]... and {len(results) - 20} more items")
+
+    else:
+        # Fallback to JSON-like display
+        console.print_json(data=results)
+
+
+def display_success(message: str) -> None:
+    """Display success message with formatting"""
+    console.print(f"[bold green]✅ {message}")
+
+
+def display_warning(message: str) -> None:
+    """Display warning message with formatting"""
+    console.print(f"[bold yellow]⚠️ {message}")
+
+
+def display_error(message: str) -> None:
+    """Display error message with formatting"""
+    console.print(f"[bold red]❌ {message}")
+
+
+def display_info(message: str) -> None:
+    """Display info message with formatting"""
+    console.print(f"[cyan]ℹ️ {message}")
+
+
+def confirm_action(message: str, default: bool = False) -> bool:
+    """Get user confirmation for dangerous actions"""
+    return click.confirm(f"⚠️ {message}", default=default)
+
+
+def get_agent_type_enum(agent_type: str):
+    """Convert string agent type to AgentTypeEnum"""
+    from hackagent.router.types import AgentTypeEnum
+
+    # Normalize the input
+    normalized = agent_type.upper().replace("-", "_").replace(" ", "_")
+
+    # Map common variations
+    type_mapping = {
+        "GOOGLE_ADK": AgentTypeEnum.GOOGLE_ADK,
+        "GOOGLE-ADK": AgentTypeEnum.GOOGLE_ADK,
+        "ADK": AgentTypeEnum.GOOGLE_ADK,
+        "LANGCHAIN": AgentTypeEnum.LANGCHAIN,
+        "LANG_CHAIN": AgentTypeEnum.LANGCHAIN,
+        "LITELLM": AgentTypeEnum.LITELLM,
+        "LITE_LLM": AgentTypeEnum.LITELLM,
+        "OPENAI_SDK": AgentTypeEnum.OPENAI_SDK,
+        "OPENAI-SDK": AgentTypeEnum.OPENAI_SDK,
+        "OPENAI": AgentTypeEnum.OPENAI_SDK,
+        "OTHER": AgentTypeEnum.UNKNOWN,
+        "UNKNOWN": AgentTypeEnum.UNKNOWN,
+    }
+
+    if normalized in type_mapping:
+        return type_mapping[normalized]
+
+    try:
+        return AgentTypeEnum(normalized)
+    except ValueError:
+        # If the type is not recognized, fallback to UNKNOWN
+        console.print(
+            f"[yellow]⚠️ Agent type '{agent_type}' not recognized, using 'UNKNOWN'[/yellow]"
+        )
+        return AgentTypeEnum.UNKNOWN
+
+
+def format_duration(seconds: float) -> str:
+    """Format duration in seconds to human readable format"""
+    if seconds < 60:
+        return f"{seconds:.1f}s"
+    elif seconds < 3600:
+        minutes = seconds / 60
+        return f"{minutes:.1f}m"
+    else:
+        hours = seconds / 3600
+        return f"{hours:.1f}h"
+
+
+def create_status_panel(title: str, content: str, status: str = "info") -> Panel:
+    """Create a status panel with appropriate styling"""
+    style_map = {
+        "success": "green",
+        "error": "red",
+        "warning": "yellow",
+        "info": "cyan",
+    }
+
+    style = style_map.get(status, "cyan")
+    return Panel(
+        Text(content, style=style), title=title, border_style=style, padding=(1, 2)
+    )
+
+
+def launch_tui(cli_config, initial_tab: str = "dashboard", initial_data: dict = None):
+    """Launch the TUI application with specified tab and optional initial data
+
+    Args:
+        cli_config: CLI configuration object
+        initial_tab: Which tab to show initially (default: "dashboard")
+        initial_data: Initial data to pre-fill in the tab (default: None)
+    """
+    try:
+        from hackagent.cli.tui import HackAgentTUI
+
+        app = HackAgentTUI(
+            cli_config, initial_tab=initial_tab, initial_data=initial_data
+        )
+        app.run()
+
+    except ImportError:
+        console.print("[bold red]❌ TUI dependencies not installed[/bold red]")
+        console.print("\n[cyan]💡 Install with:[/cyan]")
+        console.print("  uv add textual")
+        raise click.ClickException("TUI dependencies not installed")
+    except Exception as e:
+        console.print(f"[bold red]❌ TUI failed to start: {e}[/bold red]")
+        raise click.ClickException(f"TUI failed to start: {e}")