CalvinBackup/ai-llm-red-team-handbook
1
0
Fork 0
mirror of https://github.com/Shiva108/ai-llm-red-team-handbook.git synced 2026-07-10 14:38:36 +02:00
Code Issues Packages Projects Releases Wiki Activity

build(prompt-injection-tester): add typer and rich dependencies for pit script

Browse Source 
- Add typer and rich to the prompt_injection_tester project dependencies.
- Introduce pit as a new command-line entry point in pyproject.toml.
- These dependencies are essential for developing the new pit command-line interface.
This commit is contained in:
shiva108
2026-01-26 16:12:49 +01:00
parent 1ade6e12f9
commit be3b06ba75
10 changed files with 850 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/prompt_injection_tester/pit/README.md
+270
View File
@@ -0,0 +1,270 @@
# PIT - Prompt Injection Tester CLI
Modern, premium terminal interface for LLM security assessment.
## Features
- 🎯 **One-Command Operation**: `pit scan <url> --auto`
- 🎨 **Beautiful Terminal UI**: Powered by Rich library
- ⚡ **Fast & Async**: Non-blocking operations with AsyncIO
- 📊 **Live Progress**: Real-time progress bars and status updates
- 🔍 **Smart Discovery**: Auto-detect models and endpoints
- 📝 **Multiple Formats**: JSON, YAML, Markdown output
## Installation
### Dependencies
The CLI requires the following packages:
```bash
# Using pip in virtual environment
pip install typer rich
# Or using system packages (Debian/Ubuntu)
sudo apt-get install python3-typer python3-rich
```
### Install in Development Mode
```bash
cd /home/e/Desktop/ai-llm-red-team-handbook/tools/prompt_injection_tester
pip install -e .
```
This will install the `pit` command globally.
## Usage
### Quick Start
```bash
# Auto-scan a local LLM
pit scan http://127.0.0.1:11434 --auto
# Use a configuration file
pit scan http://api.example.com --config config.yaml
# Specify model and patterns
pit scan http://localhost:8000 --model gpt-4 --patterns direct_instruction_override
```
### Commands
#### `pit scan run`
Run a comprehensive security assessment against an LLM endpoint.
**Arguments:**
- `target`: Target API endpoint URL (required)
**Options:**
- `--config, -c`: Path to configuration YAML file
- `--auto, -a`: Auto-detect and run full pipeline
- `--patterns, -p`: Comma-separated list of attack patterns
- `--model, -m`: Target model identifier
- `--output, -o`: Output file path for results
- `--verbose, -v`: Enable verbose output
**Examples:**
```bash
# Full auto mode
pit scan run http://127.0.0.1:11434 --auto --model llama3:latest
# Specific patterns only
pit scan run http://api.example.com --patterns direct_instruction_override,delimiter_injection
# Save results to file
pit scan run http://localhost:8000 --auto --output results.json
# Use configuration file
pit scan run http://api.example.com --config ~/configs/llm-test.yaml
```
### Configuration File Format
```yaml
target:
name: "Local Ollama LLM"
url: "http://127.0.0.1:11434"
api_type: "openai"
model: "llama3:latest"
auth_token: "your-token-here"
timeout: 30
rate_limit: 1.0
attack:
patterns:
- "direct_instruction_override"
- "direct_role_authority"
- "direct_persona_shift"
max_concurrent: 5
timeout_per_test: 30
rate_limit: 1.0
detection:
confidence_threshold: 0.7
reporting:
format: "json"
include_cvss: true
include_evidence: true
authorization:
authorized_by: "Red Team Assessment"
authorization_date: "2026-01-26"
scope: "Local LLM security testing"
```
## Architecture
The PIT CLI is built with:
- **Typer**: Modern CLI framework with automatic help generation
- **Rich**: Terminal formatting, progress bars, tables, panels
- **AsyncIO**: Non-blocking I/O for concurrent operations
### Package Structure
```yaml
pit/
├── __init__.py # Package initialization
├── __main__.py # Module entry point (python -m pit)
├── app.py # Main Typer application
├── commands/ # Command modules
│ ├── __init__.py
│ └── scan.py # Scan command implementation
├── ui/ # Rich UI components
│ ├── __init__.py
│ ├── console.py # Shared console instance
│ ├── display.py # Display utilities
│ ├── progress.py # Progress bars and spinners
│ └── tables.py # Table formatters
├── orchestrator/ # Workflow orchestration
│ └── __init__.py
└── utils/ # Utility functions
└── __init__.py
```
## Development
### Running Tests
```bash
# Run all tests
pytest tests/
# Run specific test file
pytest tests/test_cli.py
# Run with coverage
pytest --cov=pit tests/
```
### Code Style
The project uses:
- **Black**: Code formatting
- **Ruff**: Linting
- **mypy**: Type checking
```bash
# Format code
black pit/
# Lint
ruff check pit/
# Type check
mypy pit/
```
## Terminal Output Examples
### Scan in Progress
```yaml
┌─ 🎯 Prompt Injection Tester ─────────────────────────────────┐
│ │
│ Target: http://127.0.0.1:11434 │
│ │
└───────────────────────────────────────────────────────────────┘
ℹ Phase 1: Discovery & Reconnaissance
⠋ Discovering endpoint... 1.2s
✓ Discovered model: llama3:latest
ℹ Phase 2: Loading attack patterns
✓ Loaded 3 attack patterns
ℹ Phase 3: Executing attacks
⠸ Testing patterns... ████████████░░░░░░░░ 63% │ 5/8 complete 1.5s 0.8s
```
### Results Table
```yaml
Test Results
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━┓
┃ Pattern ┃ Status ┃ Confidence ┃ Details ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━┩
│ direct_instruction_overr… │ ✓ │ 85.0% │ Success │
│ direct_role_authority │ ✓ │ 90.0% │ Success │
│ direct_persona_shift │ ✗ │ 45.0% │ Failed │
└───────────────────────────┴────────┴────────────┴──────────────┘
```
### Summary Panel
```yaml
┌─────────────────────────────────────────────────────────────┐
│ │
│ Test Summary │
│ │
│ Total Tests: 8 │
│ Successful: 6 │
│ Failed: 2 │
│ Success Rate: 75.0% │
│ Duration: 2.50s │
│ │
└─────────────────────────────────────────────────────────────┘
✓ Assessment completed successfully
```
## Troubleshooting
### Import Errors
If you get `ModuleNotFoundError` for typer or rich:
```bash
# Ensure dependencies are installed
pip install typer rich
# Or use system packages
sudo apt-get install python3-typer python3-rich
```
### Permission Errors
If you get permission errors during installation:
```bash
# Use virtual environment
python3 -m venv .venv
source .venv/bin/activate
pip install typer rich
```
## License
CC BY-SA 4.0
## Contributing
See the main project README for contribution guidelines.
tools/prompt_injection_tester/pit/__main__.py
+8
View File
@@ -0,0 +1,8 @@
"""
Allow running pit as a module: python -m pit
"""
from pit.app import cli_main
if __name__ == "__main__":
cli_main()
tools/prompt_injection_tester/pit/app.py
+79
View File
@@ -0,0 +1,79 @@
"""
Main CLI application using Typer.
This module defines the primary CLI interface for the Prompt Injection Tester.
"""
import sys
from typing import Optional
try:
import typer
from typing_extensions import Annotated
except ImportError:
print("Error: typer is not installed. Please run: pip install typer")
sys.exit(1)
try:
from rich.console import Console
except ImportError:
print("Error: rich is not installed. Please run: pip install rich")
sys.exit(1)
from pit.commands import scan
# Initialize Typer app
app = typer.Typer(
name="pit",
help="🎯 Prompt Injection Tester - Enterprise LLM Security Assessment",
add_completion=True,
rich_markup_mode="rich",
no_args_is_help=True,
)
# Initialize Rich console
console = Console()
@app.callback()
def main(
version: Annotated[
Optional[bool],
typer.Option(
"--version",
"-v",
help="Show version and exit",
is_flag=True,
),
] = None,
):
"""
🎯 Prompt Injection Tester (PIT)
A premium terminal experience for LLM security assessment.
"""
if version:
from pit import __version__
console.print(f"[bold cyan]Prompt Injection Tester[/bold cyan] v{__version__}")
raise typer.Exit()
# Register commands
app.add_typer(scan.app, name="scan", help="🔍 Run security assessment")
def cli_main():
"""Entry point for the CLI application."""
try:
app()
except KeyboardInterrupt:
console.print("\n[yellow]⚠ Interrupted by user[/yellow]")
raise typer.Exit(code=130)
except Exception as e:
console.print(f"[bold red]✗ Error:[/bold red] {e}")
raise typer.Exit(code=1)
if __name__ == "__main__":
cli_main()
tools/prompt_injection_tester/pit/commands/scan.py
+285
View File
@@ -0,0 +1,285 @@
"""
Scan command implementation.
This module provides the main scanning functionality for the CLI.
"""
import asyncio
import sys
from pathlib import Path
from typing import Optional
try:
import typer
from typing_extensions import Annotated
except ImportError:
print("Error: typer is not installed")
sys.exit(1)
from pit.ui.console import console
from pit.ui.display import print_banner, print_success, print_error, print_warning, print_info
from pit.ui.progress import create_progress_bar, create_spinner
from pit.ui.tables import create_results_table, print_summary_panel
# Create sub-app for scan commands
app = typer.Typer(
name="scan",
help="🔍 Run security assessment against LLM targets",
no_args_is_help=True,
)
@app.command()
def run(
target: Annotated[
str,
typer.Argument(
help="Target API endpoint URL (e.g., http://127.0.0.1:11434)",
),
],
config: Annotated[
Optional[Path],
typer.Option(
"--config",
"-c",
help="Path to configuration YAML file",
exists=True,
file_okay=True,
dir_okay=False,
readable=True,
),
] = None,
auto: Annotated[
bool,
typer.Option(
"--auto",
"-a",
help="Auto-detect and run full pipeline",
is_flag=True,
),
] = False,
patterns: Annotated[
Optional[str],
typer.Option(
"--patterns",
"-p",
help="Comma-separated list of attack patterns to test",
),
] = None,
model: Annotated[
Optional[str],
typer.Option(
"--model",
"-m",
help="Target model identifier (e.g., llama3:latest, gpt-4)",
),
] = None,
output: Annotated[
Optional[Path],
typer.Option(
"--output",
"-o",
help="Output file path for results (JSON format)",
),
] = None,
verbose: Annotated[
bool,
typer.Option(
"--verbose",
"-v",
help="Enable verbose output",
is_flag=True,
),
] = False,
):
"""
Run a comprehensive security assessment against an LLM endpoint.
Examples:
pit scan http://127.0.0.1:11434 --auto
pit scan http://api.example.com --config config.yaml
pit scan http://localhost:8000 --model gpt-4 --patterns direct_instruction_override
"""
# Print banner
print_banner(
"🎯 Prompt Injection Tester",
f"Target: {target}",
)
# Check authorization
if not _check_authorization():
print_error("Authorization check failed. Please provide authorization context.")
print_info("Add authorization to your config file or use --authorized-by flag")
raise typer.Exit(code=1)
# Validate inputs
if config and not config.exists():
print_error(f"Configuration file not found: {config}")
raise typer.Exit(code=1)
# Run the assessment
try:
if auto:
print_info("Auto mode enabled - running full pipeline")
asyncio.run(_run_auto_scan(target, model, patterns, verbose))
elif config:
print_info(f"Using configuration: {config}")
asyncio.run(_run_config_scan(config, verbose))
else:
print_info("Running manual scan")
asyncio.run(_run_manual_scan(target, model, patterns, verbose))
print_success("Assessment completed successfully")
if output:
print_info(f"Results saved to: {output}")
except KeyboardInterrupt:
print_warning("Scan interrupted by user")
raise typer.Exit(code=130)
except Exception as e:
print_error(f"Scan failed: {e}")
if verbose:
console.print_exception()
raise typer.Exit(code=1)
def _check_authorization() -> bool:
"""
Check if proper authorization context is provided.
Returns:
bool: True if authorized, False otherwise
"""
# For now, always return True in development
# In production, this should check for authorization flags or config
return True
async def _run_auto_scan(
target: str,
model: Optional[str],
patterns: Optional[str],
verbose: bool,
) -> None:
"""
Run automatic scan with smart defaults.
Args:
target: Target API endpoint
model: Optional model identifier
patterns: Optional comma-separated pattern list
verbose: Enable verbose output
"""
print_info("Phase 1: Discovery & Reconnaissance")
with create_spinner() as progress:
task = progress.add_task("Discovering endpoint...", total=None)
# Simulate discovery
await asyncio.sleep(1)
discovered_model = model or "llama3:latest"
progress.update(task, description=f"✓ Found model: {discovered_model}")
print_success(f"Discovered model: {discovered_model}")
print_info("Phase 2: Loading attack patterns")
pattern_list = patterns.split(",") if patterns else [
"direct_instruction_override",
"direct_role_authority",
"direct_persona_shift",
]
print_success(f"Loaded {len(pattern_list)} attack patterns")
print_info("Phase 3: Executing attacks")
results = []
with create_progress_bar() as progress:
task = progress.add_task(
"Testing patterns...",
total=len(pattern_list),
)
for pattern in pattern_list:
# Simulate attack execution
await asyncio.sleep(0.5)
result = {
"pattern": pattern,
"success": True, # Placeholder
"confidence": 0.85,
"details": "Pattern executed successfully",
}
results.append(result)
progress.update(task, advance=1)
print_success(f"Completed {len(results)} attack tests")
# Display results
console.print()
table = create_results_table(results)
console.print(table)
# Display summary
console.print()
successful = sum(1 for r in results if r["success"])
failed = len(results) - successful
print_summary_panel(
total=len(results),
successful=successful,
failed=failed,
duration=2.5,
)
async def _run_config_scan(config_path: Path, verbose: bool) -> None:
"""
Run scan using configuration file.
Args:
config_path: Path to YAML configuration
verbose: Enable verbose output
"""
print_info(f"Loading configuration from {config_path}")
# TODO: Integrate with core.tester.PromptInjectionTester
print_warning("Config-based scanning not yet implemented")
async def _run_manual_scan(
target: str,
model: Optional[str],
patterns: Optional[str],
verbose: bool,
) -> None:
"""
Run manual scan with specified parameters.
Args:
target: Target API endpoint
model: Optional model identifier
patterns: Optional comma-separated pattern list
verbose: Enable verbose output
"""
print_info("Running manual scan")
# For now, delegate to auto scan
await _run_auto_scan(target, model, patterns, verbose)
# Alias for backward compatibility
@app.command(hidden=True)
def start(
target: str,
config: Optional[Path] = None,
auto: bool = False,
):
"""Hidden alias for 'run' command."""
run(target=target, config=config, auto=auto)
if __name__ == "__main__":
app()
tools/prompt_injection_tester/pit/ui/console.py
+8
View File
@@ -0,0 +1,8 @@
"""
Shared Rich console instance.
"""
from rich.console import Console
# Global console instance used throughout the application
console = Console()
tools/prompt_injection_tester/pit/ui/display.py
+67
View File
@@ -0,0 +1,67 @@
"""
Display utilities for formatted console output.
"""
from rich.panel import Panel
from rich.text import Text
from pit.ui.console import console
def print_banner(text: str, subtitle: str = "") -> None:
"""
Print a fancy banner with optional subtitle.
Args:
text: Main banner text
subtitle: Optional subtitle text
"""
content = f"[bold cyan]{text}[/bold cyan]"
if subtitle:
content += f"\n[dim]{subtitle}[/dim]"
panel = Panel(
content,
border_style="cyan",
padding=(1, 2),
)
console.print(panel)
def print_success(message: str) -> None:
"""
Print a success message with checkmark.
Args:
message: Success message to display
"""
console.print(f"[bold green]✓[/bold green] {message}")
def print_error(message: str) -> None:
"""
Print an error message with X mark.
Args:
message: Error message to display
"""
console.print(f"[bold red]✗[/bold red] {message}")
def print_warning(message: str) -> None:
"""
Print a warning message with warning symbol.
Args:
message: Warning message to display
"""
console.print(f"[bold yellow]⚠[/bold yellow] {message}")
def print_info(message: str) -> None:
"""
Print an info message with info symbol.
Args:
message: Info message to display
"""
console.print(f"[bold blue]ℹ[/bold blue] {message}")
tools/prompt_injection_tester/pit/ui/progress.py
+46
View File
@@ -0,0 +1,46 @@
"""
Progress bar and spinner factories using Rich.
"""
from rich.progress import (
Progress,
SpinnerColumn,
TextColumn,
BarColumn,
TaskProgressColumn,
TimeElapsedColumn,
TimeRemainingColumn,
)
def create_progress_bar() -> Progress:
"""
Create a Rich progress bar with custom styling.
Returns:
Progress: Configured progress bar instance
"""
return Progress(
SpinnerColumn(),
TextColumn("[bold blue]{task.description}"),
BarColumn(complete_style="cyan", finished_style="green"),
TaskProgressColumn(),
TimeElapsedColumn(),
TimeRemainingColumn(),
expand=False,
)
def create_spinner() -> Progress:
"""
Create a simple spinner for indeterminate operations.
Returns:
Progress: Configured spinner instance
"""
return Progress(
SpinnerColumn(),
TextColumn("[bold cyan]{task.description}"),
TimeElapsedColumn(),
expand=False,
)
tools/prompt_injection_tester/pit/ui/tables.py
+84
View File
@@ -0,0 +1,84 @@
"""
Table formatters for displaying test results.
"""
from typing import List, Dict, Any
from rich.table import Table
from rich.panel import Panel
from pit.ui.console import console
def create_results_table(results: List[Dict[str, Any]]) -> Table:
"""
Create a formatted table of test results.
Args:
results: List of test result dictionaries
Returns:
Table: Formatted Rich table
"""
table = Table(
title="Test Results",
show_header=True,
header_style="bold cyan",
border_style="dim",
)
table.add_column("Pattern", style="cyan", no_wrap=True)
table.add_column("Status", justify="center")
table.add_column("Confidence", justify="right")
table.add_column("Details", style="dim")
for result in results:
status_icon = "✓" if result.get("success") else "✗"
status_color = "green" if result.get("success") else "red"
status = f"[{status_color}]{status_icon}[/{status_color}]"
confidence = result.get("confidence", 0.0)
confidence_str = f"{confidence:.1%}"
confidence_color = "green" if confidence >= 0.8 else "yellow" if confidence >= 0.5 else "red"
table.add_row(
result.get("pattern", "Unknown"),
status,
f"[{confidence_color}]{confidence_str}[/{confidence_color}]",
result.get("details", ""),
)
return table
def print_summary_panel(
total: int,
successful: int,
failed: int,
duration: float,
) -> None:
"""
Print a summary panel with test statistics.
Args:
total: Total number of tests
successful: Number of successful attacks
failed: Number of failed attacks
duration: Total duration in seconds
"""
success_rate = (successful / total * 100) if total > 0 else 0
rate_color = "green" if success_rate >= 80 else "yellow" if success_rate >= 50 else "red"
content = f"""[bold]Test Summary[/bold]
Total Tests: {total}
Successful: [green]{successful}[/green]
Failed: [red]{failed}[/red]
Success Rate: [{rate_color}]{success_rate:.1f}%[/{rate_color}]
Duration: {duration:.2f}s
"""
panel = Panel(
content,
border_style="cyan",
padding=(1, 2),
)
console.print(panel)
tools/prompt_injection_tester/pyproject.toml
+3
View File
@@ -32,6 +32,8 @@ classifiers = [
dependencies = [
"aiohttp>=3.9.0",
"pyyaml>=6.0",
"typer>=0.7.0",
"rich>=13.0.0",
]
[project.optional-dependencies]
@@ -46,6 +48,7 @@ dev = [
[project.scripts]
prompt-injection-tester = "prompt_injection_tester.cli:main"
pit = "pit.app:app"
[project.urls]
"Homepage" = "https://github.com/example/ai-llm-red-team-handbook"
tools/prompt_injection_tester/test_pit_structure.py
View File