prompt-injection-lab

A small, target-agnostic harness for testing whether an agent or user-facing LLM app can be exploited via prompt injection and jailbreaks — and a guided way to learn how those attacks work.

This is a defensive red-teaming tool: point it at a target you own (or are authorized to test), it fires a battery of attack prompts, and it reports which vulnerabilities are present. Operating assumption: any model — frontier or not — may be injectable/jailbreakable, so it has to earn its score by resisting. Only test systems you control.

SUMMARY: 47/51 attacks succeeded  (92% exploit rate)   # against the naive demo target
          severity: critical 7, high 22, medium 16, low 2

Managed with uv. The default LLM provider is Gemini (google-genai); the active provider/model is switchable via .env.

cd prompt-injection-lab
uv sync                       # creates .venv and installs deps (google-genai, dotenv)
cp .env.example .env          # then put your GEMINI_API_KEY in .env

Optional extras for other providers: uv sync --extra anthropic, --extra openai, --extra http, or --extra all.

It also ships with a deliberately-vulnerable simulated agent so you can see the harness work with zero keys:

uv run python -m pi_lab.cli list                  # see the attack library
uv run python -m pi_lab.cli run --target mock -v  # run all attacks, show reasons + replies
uv run python examples/demo.py                    # programmatic API + a custom attack

The CLI exits non-zero (code 2) when any attack succeeds — drop it into CI as a guardrail check.

With GEMINI_API_KEY set in .env, the default target is Gemini, so just:

uv run python -m pi_lab.cli run -v                      # -> gemini-3.5-flash
uv run python -m pi_lab.cli run --model gemini-3.1-pro  # pick another Gemini model

# .env
PI_LAB_PROVIDER=gemini      # gemini (default) | anthropic | openai | http | mock
# PI_LAB_MODEL=gemini-3.5-flash   # optional; else the provider default is used
GEMINI_API_KEY=...

…or override per-run on the CLI (flags beat .env):

uv run python -m pi_lab.cli run --target anthropic --model claude-opus-4-8 -v
uv run python -m pi_lab.cli run --target openai --model gpt-4o-mini

# YOUR OWN agent, via its HTTP endpoint
uv run python -m pi_lab.cli run --target http \
    --url http://localhost:8000/chat \
    --prompt-field message --response-path reply \
    --md report.md --json report.json

Precedence everywhere: CLI flag → .env → built-in default (gemini).

Beyond the fixed list, an attacker LLM can craft an attack, read how the target defended, and refine — looping until it breaks through. Faithful to PAIR + TAP (multi-stream breadth, pruning, strategy-diverse streams):

uv run python -m pi_lab.cli adaptive --target gemini --streams 3 --iterations 8 -v

It targets only the fake planted secret (no real-harm objective). Add --multi-turn for an automated Crescendo (one escalating conversation per stream instead of fresh single-shot attempts):

uv run python -m pi_lab.cli adaptive --target gemini --multi-turn --max-turns 6 -v

⚠️ Frontier models refuse to be the attacker. Gemini/Claude/GPT will decline to generate attacks — so the harness fires curated seed attacks as a fallback and flags the refusal rate. For true adaptive refinement, point --attacker at an uncensored/open model, e.g. a local Ollama: --attacker openai --base-url http://localhost:11434/v1 --model llama3.

A tool-using agent may refuse to print a secret yet still be tricked into using it in an action (an email body, a URL). ToolAgentTarget is a Gemini agent with mock tools (send_email, fetch_url, …) that records every tool call, so the harness can catch a secret leaving via an action — including indirect injection from a poisoned page/note the agent reads:

uv run python examples/tool_abuse.py     # -> toolabuse-report.json

The flagship: a single command that runs the three high-severity agent scenario families against a tool-using agent and decides pass/fail from the actions, not just the text — secret/prompt leak, tool-abuse/confused-deputy, and insecure output handling (SQLi · XSS · markdown-image exfil · SSRF · RCE). Every finding ships a fix:

uv run python -m pi_lab.cli audit -v                 # vs a hardened agent
uv run python -m pi_lab.cli audit --weak             # vs a careless config (the contrast)
uv run python -m pi_lab.cli audit --json out.json --sarif out.sarif   # CI artifacts

Oracles judge the action (the tool-call arg), not the model's chatter — so a model that refuses by quoting curl|bash is correctly read as safe.

Point it at your system — no rewrite needed:

from pi_lab import CallableTarget          # wrap any fn(prompt, system) -> str
target = CallableTarget(my_agent.chat)

…or via the CLI: --target http --url http://localhost:8000/chat for an HTTP app.

Attacks, targets, and oracles are all pluggable — a plugin is a .py that registers them:

from pi_lab import register_target, register_oracle, register_attacks, CallableTarget
register_target("myapp")(lambda **kw: CallableTarget(my_agent.chat, name="myapp"))

uv run python -m pi_lab.cli --plugin myplugin.py run --target myapp

See examples/plugin_example.py (runs offline).

Agentic-native mode — invoke the red-team conversationally from Claude Code / Cursor / any MCP client ("scan my agent for tool-abuse"):

uv sync --extra mcp
uv run python -m pi_lab.mcp_server         # exposes list_attacks / run_redteam / audit_agent

CLI · library/embed · MCP server · pre-commit hook · CI gate (exit-code + SARIF) · shareable HTML report (audit --html report.html) · Excel/JSON · vendor/model comparison. run/audit exit 2 when anything gets through. Full menu with copy-paste examples: docs/05-usage-modes.md.

Roll up every run (static vs Gemini + mock, adaptive, multi-turn, tool-abuse, agent audit) into one colour-coded workbook — prompts, responses, WORKED/REFUSED:

uv run python examples/export_excel.py   # -> results.xlsx

51 attacks across 10 categories — a starter battery that grows. The invisible / obfuscation_advanced ones are generated from pi_lab/encoders.py (invisible Unicode can't live in JSON); the rest are JSON you can copy and extend. See docs/02-research.md for the technique catalogue and sources.

Attacks are plain JSON in data/attacks/ — add your own by copying a case and giving it a unique id + a success spec.

attacks (JSON)  ──►  Runner  ──►  Target.send(prompt, system)  ──►  Judge  ──►  Report
                       │            mock | anthropic | openai | http      │
              plants a secret +                                    canary / leaked-secret /
              rules in the system prompt                           refusal / LLM-as-judge

• Target (pi_lab/targets.py) — the one abstraction that makes it target-agnostic: gemini (default), anthropic, openai, http, mock. Implement send() to test anything.
• Config (pi_lab/config.py) — loads .env, resolves the default provider (gemini) and model.
• Judge (pi_lab/judges.py) — HeuristicJudge (canaries/secret/refusal, default) or LLMJudge (model grades open-ended jailbreaks).
• Runner (pi_lab/runner.py) — owns the test scenario (a system prompt with a secret to protect and rules to keep) and runs every attack.

• docs/01-concepts.md — what each attack class is, with the why.
• docs/02-research.md — deeper technique catalogue, layered defenses, and pointers to OWASP LLM Top 10, MITRE ATLAS, and bigger tools (Garak, PyRIT, promptfoo) when you outgrow this.

prompt-injection-lab/
├── pyproject.toml     # uv-managed deps (google-genai default + optional extras)
├── .env.example       # copy to .env: provider, model, API keys
├── BUILD_ORDER.md     # what's built vs pending (living tracker)
├── pi_lab/            # harness core
│   ├── targets.py     # Target adapters (mock/gemini/…/http/tool_agent/callable)
│   ├── attacks.py · generators.py · encoders.py   # the attack library + smuggling
│   ├── judges.py · oracles.py     # verdicts (text) + Oracles (text + side-effects)
│   ├── adaptive.py    # PAIR/TAP single-turn + Crescendo multi-turn
│   ├── audit.py       # Attack × Target × Oracle agent audit (the "do" layer)
│   ├── registry.py    # plugin API (register targets/oracles/attacks)
│   └── runner.py · report.py · cli.py · config.py
├── data/attacks/      # the attack library (JSON — edit/extend freely)
├── docs/              # 01-concepts · 02-research · 03-coverage · 04-positioning-and-roadmap
└── examples/          # demo · probe · tool_abuse · export_excel · plugin_example

This is a focused tool — strong on injection/jailbreak (the "say" layer) and the agent "do" layer (tool-abuse, insecure output handling). It is not exhaustive: see docs/03-coverage.md for the honest map of what's tested vs not, and pair it with garak / promptfoo / PyRIT / MCP scanners for the rest.

Defensive tool — test only systems you own or are authorized to test. It is benign by construction (planted fake secret + canary tokens, no harmful-content library). See RESPONSIBLE_USE.md. Licensed under Apache-2.0 (LICENSE).