From 9ed5b09affe64e3840a0d419b3d624f854ffdc04 Mon Sep 17 00:00:00 2001 From: shiva108 Date: Mon, 26 Jan 2026 20:30:28 +0100 Subject: [PATCH] feat: Introduce the Prompt Injection Tester (PIT) tool, including its CI workflow, Docker setup, and comprehensive changelog. --- .github/workflows/pit-test.yml | 184 ++++++++++++++ .gitignore | 11 + tools/prompt_injection_tester/.dockerignore | 85 +++++++ tools/prompt_injection_tester/CHANGELOG.md | 236 ++++++++++++++++++ tools/prompt_injection_tester/Dockerfile | 49 ++++ .../docker-compose.yml | 87 +++++++ 6 files changed, 652 insertions(+) create mode 100644 .github/workflows/pit-test.yml create mode 100644 tools/prompt_injection_tester/.dockerignore create mode 100644 tools/prompt_injection_tester/CHANGELOG.md create mode 100644 tools/prompt_injection_tester/Dockerfile create mode 100644 tools/prompt_injection_tester/docker-compose.yml diff --git a/.github/workflows/pit-test.yml b/.github/workflows/pit-test.yml new file mode 100644 index 0000000..f855acd --- /dev/null +++ b/.github/workflows/pit-test.yml @@ -0,0 +1,184 @@ +name: PIT - Continuous Integration + +on: + push: + branches: [main, develop] + paths: + - 'tools/prompt_injection_tester/**' + - '.github/workflows/pit-test.yml' + pull_request: + branches: [main, develop] + paths: + - 'tools/prompt_injection_tester/**' + workflow_dispatch: + +jobs: + lint-and-test: + name: Lint and Test + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ['3.10', '3.11', '3.12'] + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: 'pip' + + - name: Install dependencies + working-directory: tools/prompt_injection_tester + run: | + python -m pip install --upgrade pip + pip install -e . + pip install pytest pytest-cov pytest-asyncio black ruff mypy + + - name: Lint with ruff + working-directory: tools/prompt_injection_tester + run: | + ruff check pit/ tests/ --output-format=github + continue-on-error: true + + - name: Format check with black + working-directory: tools/prompt_injection_tester + run: | + black --check pit/ tests/ + continue-on-error: true + + - name: Type check with mypy + working-directory: tools/prompt_injection_tester + run: | + mypy pit/ --ignore-missing-imports + continue-on-error: true + + - name: Run integration tests + working-directory: tools/prompt_injection_tester + run: | + pytest tests/integration/ -v --cov=pit --cov-report=xml --cov-report=term + + - name: Run report formatter tests + working-directory: tools/prompt_injection_tester + run: | + python tests/test_reports.py + + - name: Upload coverage reports + uses: codecov/codecov-action@v4 + with: + file: tools/prompt_injection_tester/coverage.xml + flags: unittests + name: codecov-${{ matrix.python-version }} + continue-on-error: true + + docker-build: + name: Docker Build Test + runs-on: ubuntu-latest + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Build Docker image + working-directory: tools/prompt_injection_tester + run: | + docker build -t pit:test . + + - name: Test Docker image + run: | + docker run --rm pit:test --version + docker run --rm pit:test --help + + e2e-test: + name: End-to-End Test + runs-on: ubuntu-latest + needs: [lint-and-test] + + services: + ollama: + image: ollama/ollama:latest + ports: + - 11434:11434 + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + cache: 'pip' + + - name: Install dependencies + working-directory: tools/prompt_injection_tester + run: | + python -m pip install --upgrade pip + pip install -e . + + - name: Wait for Ollama + run: | + timeout 60 bash -c 'until curl -f http://localhost:11434/api/tags; do sleep 2; done' + + - name: Pull Ollama model + run: | + docker exec ${{ job.services.ollama.id }} ollama pull llama3.2:1b + + - name: Run E2E tests + working-directory: tools/prompt_injection_tester + run: | + python tests/e2e_test.py --target http://localhost:11434/api/chat --model llama3.2:1b --quick + timeout-minutes: 10 + + - name: Upload test reports + if: always() + uses: actions/upload-artifact@v4 + with: + name: e2e-reports + path: tools/prompt_injection_tester/pit_report_*.* + retention-days: 7 + + security-scan: + name: Security Scan + runs-on: ubuntu-latest + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install dependencies + working-directory: tools/prompt_injection_tester + run: | + python -m pip install --upgrade pip + pip install -e . + pip install safety bandit + + - name: Run safety check + working-directory: tools/prompt_injection_tester + run: | + safety check --json + continue-on-error: true + + - name: Run bandit security scan + working-directory: tools/prompt_injection_tester + run: | + bandit -r pit/ -f json -o bandit-report.json + continue-on-error: true + + - name: Upload security reports + if: always() + uses: actions/upload-artifact@v4 + with: + name: security-reports + path: tools/prompt_injection_tester/*-report.json + retention-days: 30 diff --git a/.gitignore b/.gitignore index 2badb01..979f0f3 100644 --- a/.gitignore +++ b/.gitignore @@ -19,6 +19,7 @@ build/ .venvs/ /tools/prompt_injection_tester/.venv + # --- Node.js --- node_modules/ npm-debug.log @@ -86,8 +87,18 @@ tools/prompt_injection_tester/report.json tools/prompt_injection_tester/report.yaml tools/prompt_injection_tester/report.html tools/prompt_injection_tester/interrupted_*.json +tools/prompt_injection_tester/results.json +tools/prompt_injection_tester/results.html +tools/prompt_injection_tester/results.yaml +tools/prompt_injection_tester/config.yaml +tools/prompt_injection_tester/IMPLEMENTATION_COMPLETE.md # Legacy/Archive files tools/prompt_injection_tester/docs/reports/archive/CODE_REVIEW_2026_01_26.md tools/prompt_injection_tester/docs/specs/CORE_ARCHITECTURE_legacy.md tools/prompt_injection_tester/docs/specs/FUNCTIONAL_SPEC_v2_legacy.md +.idea +.gitignore +tools/prompt_injection_tester/RELEASE.md +.ripgreprc +tools/prompt_injection_tester/PHASE5_COMPLETE.md diff --git a/tools/prompt_injection_tester/.dockerignore b/tools/prompt_injection_tester/.dockerignore new file mode 100644 index 0000000..ac0aa1f --- /dev/null +++ b/tools/prompt_injection_tester/.dockerignore @@ -0,0 +1,85 @@ +# Git +.git/ +.gitignore +.gitattributes + +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +*.egg-info/ +.eggs/ +dist/ +build/ +.venv/ +.venvs/ +*.egg + +# Testing +.pytest_cache/ +.coverage +htmlcov/ +.tox/ + +# IDEs +.idea/ +.vscode/ +*.sublime-project +*.sublime-workspace +*.swp +*.swo +*~ + +# Documentation (build artifacts, not source) +docs/_build/ +docs/.doctrees/ + +# Reports (generated files) +pit_report_*.json +pit_report_*.yaml +pit_report_*.html +test_report_*.json +test_report_*.yaml +test_report_*.html +report.json +report.yaml +report.html +interrupted_*.json +*.log + +# CI/CD +.github/ +.gitlab-ci.yml +.travis.yml + +# Docker +Dockerfile +docker-compose*.yml +.dockerignore + +# Development files +*.md.backup +.DS_Store +Thumbs.db + +# Temporary files +*.tmp +*.temp +*.bak + +# Legacy/Archive +docs/reports/archive/ +docs/specs/*_legacy.md + +# Keep these files +!README.md +!USER_GUIDE.md +!PATTERN_DEVELOPMENT.md +!ARCHITECTURE.md +!SPECIFICATION.md +!CHANGELOG.md +!RELEASE.md +!IMPLEMENTATION_COMPLETE.md +!PHASE2_COMPLETE.md diff --git a/tools/prompt_injection_tester/CHANGELOG.md b/tools/prompt_injection_tester/CHANGELOG.md new file mode 100644 index 0000000..1087c0d --- /dev/null +++ b/tools/prompt_injection_tester/CHANGELOG.md @@ -0,0 +1,236 @@ +# Changelog + +All notable changes to the Prompt Injection Tester (PIT) will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [2.0.0] - 2026-01-26 + +### 🎉 Major Release - Complete Re-Architecture + +This release represents a complete re-engineering of the Prompt Injection Tester with a focus on reliability, performance, and user experience. + +### Added + +#### Architecture + +- **Sequential 4-Phase Pipeline Pattern** - Eliminates tool use concurrency errors + - Phase 1: Discovery - Scan target for injection points + - Phase 2: Attack - Execute patterns sequentially + - Phase 3: Verification - Analyze responses with detection framework + - Phase 4: Reporting - Generate multi-format reports +- Proper phase ordering enforcement with explicit wait points +- Comprehensive error handling with custom exception hierarchy +- Clean resource management with try/finally blocks + +#### Configuration + +- **Type-Safe Configuration** using Pydantic v2 + - `TargetConfig` - Target endpoint configuration + - `AttackConfig` - Attack pattern selection and rate limiting + - `ReportingConfig` - Output format and path configuration +- YAML configuration file support with environment variable expansion (`${VAR}`) +- CLI argument to Pydantic model conversion +- Configuration validation with helpful error messages + +#### Reporting + +- **Multi-Format Report Generation** + - **JSON** - Clean JSON with configurable formatting + - **YAML** - Human-readable YAML output + - **HTML** - Professional HTML reports with: + - Embedded CSS (no external dependencies) + - Responsive design (desktop/mobile/print) + - Color-coded severity levels + - Summary dashboard with statistics + - Detailed test results with evidence +- Abstract `ReportFormatter` base class for extensibility +- `save_report()` convenience function with format auto-detection +- Report metadata tracking (timestamp, duration, target info) + +#### Testing + +- Integration test suite ([tests/integration/test_pipeline.py](tests/integration/test_pipeline.py)) + - Pipeline structure validation + - Phase ordering tests + - Formatter functionality tests +- End-to-end test script ([tests/e2e_test.py](tests/e2e_test.py)) + - Pipeline execution tests + - Report format validation + - Error handling verification + - Support for quick testing mode +- Report validation test suite ([tests/test_reports.py](tests/test_reports.py)) + - Mock data testing (no LLM required) + - Format-specific validation + - Element presence checks + +#### Documentation + +- **USER_GUIDE.md** (~800 lines) - Comprehensive user manual + - Installation instructions + - Quick start guide + - Complete command reference + - Configuration examples + - Troubleshooting guide + - Best practices +- **PATTERN_DEVELOPMENT.md** (~650 lines) - Pattern creation guide + - Pattern architecture overview + - Step-by-step creation guide + - Pattern types (single-turn, multi-turn, composite) + - Advanced features + - Testing patterns +- **ARCHITECTURE.md** (~1,200 lines) - Technical architecture + - Sequential Pipeline Pattern definition + - Layered architecture design + - Technology stack details + - Data flow diagrams + - Testing strategy +- **SPECIFICATION.md** (~900 lines) - Functional specification + - One-command workflow design + - CLI commands and options + - UI mockups with ASCII art + - Error handling specifications + - Output schemas +- **RELEASE.md** - Deployment and release guide +- **CHANGELOG.md** - This file + +#### Deployment + +- **Dockerfile** - Multi-stage Docker build + - Based on Python 3.11-slim + - Optimized layer caching + - Health check included + - Volume support for reports +- **docker-compose.yml** - Complete development environment + - PIT service + - Ollama LLM service for testing + - Nginx for report viewing + - Network configuration + - Volume management +- **GitHub Actions CI/CD** ([.github/workflows/pit-test.yml](../../.github/workflows/pit-test.yml)) + - Lint and test (Python 3.10, 3.11, 3.12) + - Docker build test + - End-to-end testing with Ollama + - Security scanning (Safety, Bandit) + - Coverage reporting +- **.dockerignore** - Optimized Docker context + +#### CLI + +- Modern CLI interface using Typer +- Rich terminal UI with: + - Progress bars for long operations + - Spinners for phase execution + - Color-coded severity indicators + - Formatted tables for results + - Professional error messages +- `--auto` flag for one-command workflow +- `--config` flag for YAML configuration files +- `--patterns` flag for pattern selection +- `--output` flag with format auto-detection +- `--verbose` flag for debugging + +#### Core Implementation + +- `pit/config/schema.py` - Pydantic configuration models +- `pit/config/loader.py` - YAML config loading with env expansion +- `pit/errors/exceptions.py` - Custom exception hierarchy +- `pit/errors/handlers.py` - User-friendly error handling +- `pit/orchestrator/pipeline.py` - Sequential pipeline executor +- `pit/orchestrator/phases.py` - All 4 phase implementations with real logic +- `pit/reporting/formatters.py` - Multi-format report generation (~620 lines) +- Bridge integration in `pit/orchestrator/workflow.py` +- CLI integration in `pit/commands/scan.py` + +### Changed + +- **Breaking**: Complete re-architecture from v1.x +- **Breaking**: CLI interface updated (now using Typer instead of argparse) +- **Breaking**: Configuration format changed (now using Pydantic v2) +- **Breaking**: Import paths changed (new modular structure) +- Discovery phase now uses real `InjectionTester.discover_injection_points()` +- Attack phase executes patterns sequentially with rate limiting +- Verification phase uses actual confidence scores from detection framework +- Reporting phase generates professional multi-format reports +- Error handling provides user-friendly messages with suggestions +- Resource management ensures proper cleanup in all phases + +### Fixed + +- **Critical**: Eliminated tool use concurrency errors through sequential execution +- Proper cleanup of InjectionTester instances with try/finally blocks +- Race conditions in pattern execution +- Memory leaks from unclosed HTTP clients +- Inconsistent error messages +- Missing type hints throughout codebase + +### Dependencies + +- Added `httpx>=0.24.0` - HTTP client for new pipeline +- Added `pydantic>=2.0.0` - Type-safe configuration +- Added `jinja2>=3.1.0` - HTML template rendering +- Updated `typer>=0.9.0` - Modern CLI framework +- Updated `rich>=13.0.0` - Terminal UI enhancements +- Maintained `aiohttp>=3.9.0` - Core framework compatibility +- Maintained `pyyaml>=6.0` - Configuration file support + +### Security + +- Comprehensive security scanning in CI/CD pipeline +- Input validation using Pydantic models +- Safe YAML loading with `yaml.safe_load()` +- Environment variable expansion for sensitive data +- No hardcoded credentials or API keys +- Proper error handling to avoid information disclosure + +### Performance + +- Sequential execution prevents race conditions while maintaining efficiency +- Rate limiting respects API constraints +- Efficient resource cleanup +- Minimal memory footprint +- Optimized Docker image layers + +### Statistics + +- **Total Lines of Code**: ~5,000 lines (excluding tests) +- **Test Code**: ~950 lines +- **Documentation**: ~2,250 lines +- **New Files Created**: 21 files +- **Modified Files**: 4 files +- **Report Formats**: 3 (JSON, YAML, HTML) +- **Patterns Supported**: 20+ built-in patterns + +## [1.x] - Previous Versions + +### Legacy Implementation + +Previous versions (1.x) featured: + +- Basic CLI with argparse +- Concurrent pattern execution +- Single JSON report format +- Limited error handling +- Manual configuration + +**Migration Guide**: See [RELEASE.md](RELEASE.md#upgrading) for upgrading from v1.x to v2.0.0. + +--- + +## Versioning Policy + +- **Major version** (X.0.0): Breaking changes, API changes, architecture changes +- **Minor version** (2.X.0): New features, non-breaking changes +- **Patch version** (2.0.X): Bug fixes, security patches + +## Contributing + +See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on contributing to this project. + +--- + +**Project**: AI LLM Red Team Handbook +**Tool**: Prompt Injection Tester +**License**: CC BY-SA 4.0 +**Maintained By**: AI LLM Red Team Handbook Contributors diff --git a/tools/prompt_injection_tester/Dockerfile b/tools/prompt_injection_tester/Dockerfile new file mode 100644 index 0000000..27107d6 --- /dev/null +++ b/tools/prompt_injection_tester/Dockerfile @@ -0,0 +1,49 @@ +# Prompt Injection Tester (PIT) - Docker Image +# Version: 2.0.0 +# Architecture: Sequential 4-Phase Pipeline + +FROM python:3.11-slim + +# Metadata +LABEL maintainer="AI LLM Red Team Handbook" +LABEL version="2.0.0" +LABEL description="Prompt Injection Testing tool with sequential pipeline architecture" + +# Set working directory +WORKDIR /app + +# Install system dependencies +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + git \ + curl \ + ca-certificates && \ + rm -rf /var/lib/apt/lists/* + +# Copy project files +COPY pyproject.toml README.md ./ +COPY pit/ ./pit/ +COPY tests/ ./tests/ + +# Install Python dependencies +RUN pip install --no-cache-dir -e . + +# Create directory for reports +RUN mkdir -p /reports + +# Set environment variables +ENV PYTHONUNBUFFERED=1 +ENV PIT_OUTPUT_DIR=/reports + +# Health check +HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \ + CMD pit --version || exit 1 + +# Default command (show help) +ENTRYPOINT ["pit"] +CMD ["--help"] + +# Usage examples: +# Build: docker build -t pit:2.0.0 . +# Run: docker run -v $(pwd)/reports:/reports pit:2.0.0 scan http://host.docker.internal:11434/api/chat --auto +# Interactive: docker run -it pit:2.0.0 scan http://target-url --auto diff --git a/tools/prompt_injection_tester/docker-compose.yml b/tools/prompt_injection_tester/docker-compose.yml new file mode 100644 index 0000000..57462ac --- /dev/null +++ b/tools/prompt_injection_tester/docker-compose.yml @@ -0,0 +1,87 @@ +version: '3.8' + +services: + # Prompt Injection Tester + pit: + build: + context: . + dockerfile: Dockerfile + image: pit:2.0.0 + container_name: pit-scanner + volumes: + - ./reports:/reports + - ./config:/config:ro + environment: + - PIT_OUTPUT_DIR=/reports + - PYTHONUNBUFFERED=1 + # Command examples (uncomment one to use): + # command: scan http://ollama:11434/api/chat --auto --output /reports/report.html + command: --help + networks: + - pit-network + depends_on: + - ollama + + # Ollama LLM service (for testing) + ollama: + image: ollama/ollama:latest + container_name: pit-ollama + volumes: + - ollama-data:/root/.ollama + ports: + - "11434:11434" + networks: + - pit-network + healthcheck: + test: ["CMD", "curl", "-f", "http://localhost:11434/api/tags"] + interval: 30s + timeout: 10s + retries: 3 + start_period: 10s + + # Optional: Web UI for viewing HTML reports + nginx: + image: nginx:alpine + container_name: pit-reports-viewer + volumes: + - ./reports:/usr/share/nginx/html:ro + ports: + - "8080:80" + networks: + - pit-network + profiles: + - ui + +networks: + pit-network: + driver: bridge + +volumes: + ollama-data: + driver: local + +# Usage Examples: +# +# 1. Start services: +# docker-compose up -d +# +# 2. Pull Ollama model: +# docker-compose exec ollama ollama pull llama3:latest +# +# 3. Run scan (override command): +# docker-compose run --rm pit scan http://ollama:11434/api/chat --auto --model llama3:latest +# +# 4. View reports: +# docker-compose --profile ui up -d nginx +# Open http://localhost:8080 in browser +# +# 5. Run with custom config: +# docker-compose run --rm pit scan http://ollama:11434/api/chat --config /config/config.yaml +# +# 6. Quick test: +# docker-compose run --rm pit scan http://ollama:11434/api/chat --patterns direct_instruction_override +# +# 7. Generate all report formats: +# docker-compose run --rm pit scan http://ollama:11434/api/chat --auto --output /reports/report.json +# docker-compose run --rm pit scan http://ollama:11434/api/chat --auto --output /reports/report.yaml +# docker-compose run --rm pit scan http://ollama:11434/api/chat --auto --output /reports/report.html