diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..b1df615 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,31 @@ +name: CI + +on: + push: + branches: [main, master] + pull_request: + +jobs: + test: + name: Test (Python ${{ matrix.python-version }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -e ".[dev]" + + - name: Run tests + run: pytest diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..a4d0ca7 --- /dev/null +++ b/.gitignore @@ -0,0 +1,31 @@ +# Byte-compiled / optimized / cache +__pycache__/ +*.py[cod] +*$py.class + +# Distribution / packaging +build/ +dist/ +*.egg-info/ +*.egg + +# Virtual environments +.venv/ +venv/ +env/ + +# Test / coverage artifacts +.coverage +.coverage.* +htmlcov/ +.pytest_cache/ +.mypy_cache/ + +# Generated reports +*.report.json +*.report.html + +# Editor / OS +.idea/ +.vscode/ +.DS_Store diff --git a/README.md b/README.md index f607897..9779c15 100644 --- a/README.md +++ b/README.md @@ -30,7 +30,7 @@ According to OWASP's Top 10 for LLM Applications (2023), prompt injection ranks ## ✨ Features -- 🎯 **35+ High-Quality Attack Payloads** across 6 vulnerability categories +- 🎯 **36 High-Quality Attack Payloads** across 6 vulnerability categories - 🔍 **Intelligent Response Analysis** using pattern matching and heuristics - 📊 **Risk Scoring Algorithm** (0-10 scale) with severity classification - 📄 **Multiple Output Formats**: Console, JSON, and HTML reports @@ -143,7 +143,7 @@ promptscan info Target: https://api.example.com/chat Scan Date: 2026-05-09 10:00:00 UTC -Total Payloads Tested: 35 +Total Payloads Tested: 36 Vulnerabilities Found: 8 Risk Score: 8.2/10.0 @@ -172,7 +172,7 @@ implement input sanitization, and add output validation before production use. "scan_metadata": { "target": "https://api.example.com/chat", "timestamp": "2026-05-09T10:00:00Z", - "total_payloads": 35, + "total_payloads": 36, "categories_tested": 6 }, "risk_assessment": { @@ -202,15 +202,16 @@ implement input sanitization, and add output validation before production use. ### HTML Report -![HTML Report Example](docs/images/html-report-example.png) - -*Professional HTML reports with visual risk indicators and detailed vulnerability breakdowns* +Running with `-o html -r report.html` produces a self-contained styled HTML +report featuring a risk-score banner, a severity-coloured summary grid, and a +detailed vulnerabilities table. Open the generated file in any browser to view +it. --- ## 🎯 Attack Payload Categories -PromptScan includes 35+ carefully crafted payloads across 6 categories: +PromptScan includes 36 carefully crafted payloads across 6 categories: ### 1. System Prompt Leak (6 payloads) Attempts to extract the system prompt or initial instructions that define the AI's behavior. diff --git a/docs/API.md b/docs/API.md new file mode 100644 index 0000000..ad379d9 --- /dev/null +++ b/docs/API.md @@ -0,0 +1,116 @@ +# PromptScan API Documentation + +PromptScan can be used as a Python library in addition to the `promptscan` CLI. +This page documents the public API exported from the `promptscan` package. + +> Requires Python 3.9+ (see `pyproject.toml`). + +## Public objects + +```python +from promptscan import ( + AttackEngine, # Orchestrates a full scan + HTTPClient, # Sends payloads to the target API + ResponseAnalyzer, # Detects vulnerabilities in a response + RiskScorer, # Calculates risk scores and severity + ReportGenerator, # Renders console / JSON / HTML reports +) +from promptscan.config import ScanConfig +from promptscan.analyzer import VulnerabilityMatch +``` + +## Quick start (running a full scan) + +`AttackEngine.run_scan()` is asynchronous, so call it inside an event loop. + +```python +import asyncio +from pathlib import Path + +from promptscan import AttackEngine, ReportGenerator +from promptscan.config import ScanConfig + +config = ScanConfig( + target_url="https://api.example.com/chat", + api_key="YOUR_API_KEY", # optional -> sent as X-API-Key + bearer_token=None, # optional -> sent as Authorization: Bearer ... + timeout=30, + output_format="json", # "console" | "json" | "html" + output_file=Path("report.json"), # required for "json" / "html" +) + +vulnerabilities, metadata = asyncio.run(AttackEngine(config).run_scan()) + +reporter = ReportGenerator() +reporter.generate_json_report( + vulnerabilities, metadata, config.target_url, config.output_file +) +``` + +`run_scan()` returns a tuple `(vulnerabilities, metadata)`: + +- `vulnerabilities` — `list[VulnerabilityMatch]` +- `metadata` — `dict` with keys `total_payloads`, `categories_tested`, + `vulnerabilities_found`, `successful_attacks`, `failed_attacks`. If the scan + could not run (unreachable target or no payloads), `metadata` contains an + `error` key instead and `vulnerabilities` is empty — always check for this + before treating the result as a clean report. + +## `ScanConfig` + +Dataclass holding scan settings. Validates on construction and raises +`ValueError` for an empty `target_url`, an invalid `output_format`, or a +`json`/`html` format without an `output_file`. + +| Field | Type | Default | Notes | +|-------|------|---------|-------| +| `target_url` | `str` | — | Required. The endpoint that receives `POST {"message": }`. | +| `api_key` | `str \| None` | `None` | Sent as the `X-API-Key` header. | +| `bearer_token` | `str \| None` | `None` | Sent as `Authorization: Bearer ...`. | +| `timeout` | `int` | `30` | Per-request timeout in seconds. | +| `max_retries` | `int` | `3` | Retries on timeout / request error. | +| `payload_dir` | `Path \| None` | `None` | Custom payload directory; defaults to the bundled payloads. | +| `output_format` | `str` | `"console"` | `console`, `json`, or `html`. | +| `output_file` | `Path \| None` | `None` | Required for `json`/`html`. | +| `verbose` | `bool` | `False` | Verbose logging. | +| `custom_headers` | `dict[str, str] \| None` | `None` | Merged into request headers. | + +## `VulnerabilityMatch` + +Dataclass describing a single finding: + +| Field | Type | Description | +|-------|------|-------------| +| `vulnerability_type` | `str` | e.g. `"System Prompt Leak"`. | +| `severity` | `str` | `LOW`, `MEDIUM`, `HIGH`, or `CRITICAL`. | +| `confidence` | `float` | `0.0`–`1.0`. | +| `evidence` | `str` | The text snippet that triggered detection. | +| `payload_used` | `str` | The payload that was sent (truncated). | +| `response_excerpt` | `str` | Surrounding response text. | + +## Analyzing a single response + +```python +from promptscan import ResponseAnalyzer + +analyzer = ResponseAnalyzer() +response_data = {"success": True, "response_text": "You are an AI assistant ..."} +findings = analyzer.analyze_response(response_data, payload="...", payload_category="System Prompt Leak") +``` + +`response_data` must contain `success` (bool) and `response_text` (str), matching +the shape returned by `HTTPClient.send_payload()`. + +## Scoring + +```python +from promptscan import RiskScorer + +scorer = RiskScorer() +score = scorer.calculate_risk_score(findings, total_payloads=36) # 0.0 – 10.0 +severity = scorer.get_severity_level(score) # CRITICAL/HIGH/MEDIUM/LOW/MINIMAL +recommendation = scorer.get_recommendation(score) +category_scores = scorer.calculate_category_scores(findings) +``` + +See [METHODOLOGY.md](METHODOLOGY.md) for the scoring algorithm details. diff --git a/promptscan/analyzer.py b/promptscan/analyzer.py index c7b71ba..c81a1ff 100644 --- a/promptscan/analyzer.py +++ b/promptscan/analyzer.py @@ -1,5 +1,7 @@ """Response analyzer for detecting prompt injection vulnerabilities.""" +from __future__ import annotations + import re from dataclasses import dataclass from typing import Any, Dict, List, Optional diff --git a/promptscan/cli.py b/promptscan/cli.py index cce3339..fb79ab1 100644 --- a/promptscan/cli.py +++ b/promptscan/cli.py @@ -123,6 +123,15 @@ def test( engine = AttackEngine(config) vulnerabilities, metadata = asyncio.run(engine.run_scan()) + # If the scan could not run (e.g. unreachable target or no payloads), + # do NOT emit a "clean" report — that would be a misleading false + # negative for a security tool. Surface the error and exit non-zero. + if metadata.get("error"): + console.print( + f"[red]Scan did not complete: {metadata['error']}[/red]" + ) + sys.exit(2) + # Generate report reporter = ReportGenerator() @@ -201,7 +210,7 @@ def info() -> None: console.print(" vulnerabilities. Designed for security professionals and developers.\n") console.print("[bold]Features:[/bold]") - console.print(" • 25+ high-quality attack payloads across 6 categories") + console.print(" • 36 high-quality attack payloads across 6 categories") console.print(" • Intelligent response analysis and pattern matching") console.print(" • Risk scoring algorithm (0-10 scale)") console.print(" • Multiple output formats (console, JSON, HTML)") diff --git a/promptscan/engine.py b/promptscan/engine.py index 76240f8..bf43b56 100644 --- a/promptscan/engine.py +++ b/promptscan/engine.py @@ -91,7 +91,8 @@ async def run_scan(self) -> Tuple[List[VulnerabilityMatch], Dict[str, Any]]: # Execute attacks console.print("[bold yellow]Executing security tests...[/bold yellow]") all_vulnerabilities: List[VulnerabilityMatch] = [] - + flagged_payloads = 0 + with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), @@ -110,7 +111,9 @@ async def run_scan(self) -> Tuple[List[VulnerabilityMatch], Dict[str, Any]]: ) all_vulnerabilities.extend(vulnerabilities) - + if vulnerabilities: + flagged_payloads += 1 + progress.update(task, advance=1) # Small delay to avoid overwhelming the API @@ -121,8 +124,11 @@ async def run_scan(self) -> Tuple[List[VulnerabilityMatch], Dict[str, Any]]: "total_payloads": total_payloads, "categories_tested": len(payloads), "vulnerabilities_found": len(all_vulnerabilities), - "successful_attacks": len(all_vulnerabilities), - "failed_attacks": total_payloads - len(all_vulnerabilities), + # Payloads that produced at least one finding vs. those that did not. + # (A single payload can trigger multiple findings, so these are + # counted per-payload to stay consistent with total_payloads.) + "successful_attacks": flagged_payloads, + "failed_attacks": total_payloads - flagged_payloads, } console.print(f"\n[green]✓ Scan complete[/green]") diff --git a/promptscan/scorer.py b/promptscan/scorer.py index bc4e4ae..6802a0a 100644 --- a/promptscan/scorer.py +++ b/promptscan/scorer.py @@ -1,5 +1,7 @@ """Risk scoring algorithm for prompt injection vulnerabilities.""" +from __future__ import annotations + from typing import List from promptscan.analyzer import VulnerabilityMatch diff --git a/pyproject.toml b/pyproject.toml index cc9ac40..355c854 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ classifiers = [ "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", "Topic :: Security", "Topic :: Software Development :: Testing", ] diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..5f7b341 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,59 @@ +"""Tests for the command-line interface.""" + +import promptscan.cli as cli_module +from click.testing import CliRunner + +from promptscan.cli import cli + + +class _FakeEngine: + """Stand-in for AttackEngine returning a canned (vulns, metadata) tuple.""" + + result = ([], {}) + + def __init__(self, config): + self.config = config + + async def run_scan(self): + return self._result + + +def _patch_engine(monkeypatch, result): + fake = type("FakeEngine", (_FakeEngine,), {"_result": result}) + monkeypatch.setattr(cli_module, "AttackEngine", fake) + + +def test_failed_scan_exits_nonzero_without_clean_report(monkeypatch): + """A scan that cannot run must not print a misleading 'safe' report.""" + _patch_engine(monkeypatch, ([], {"error": "Connection failed"})) + + runner = CliRunner() + result = runner.invoke(cli, ["test", "http://unreachable.invalid/chat"]) + + # Must signal failure, not success. + assert result.exit_code == 2 + # Must not claim the target is well-protected. + assert "MINIMAL RISK" not in result.output + assert "Connection failed" in result.output + + +def test_successful_scan_with_no_vulnerabilities_exits_zero(monkeypatch): + """A completed scan with no findings is a genuine clean result.""" + _patch_engine( + monkeypatch, + ([], {"total_payloads": 36, "categories_tested": 6}), + ) + + runner = CliRunner() + result = runner.invoke(cli, ["test", "http://127.0.0.1/chat"]) + + assert result.exit_code == 0 + assert "PROMPTSCAN SECURITY ASSESSMENT REPORT" in result.output + + +def test_info_command_reports_correct_payload_count(): + runner = CliRunner() + result = runner.invoke(cli, ["info"]) + + assert result.exit_code == 0 + assert "36" in result.output