Skip to main content

AI security testing framework โ€” pytest for LLM applications

Project description

Vektor ๐Ÿ›ก๏ธ

pytest for AI security โ€” scan LLM apps for vulnerabilities in 30 seconds

License: MIT Python 3.8+

๐ŸŽฏ What is Vektor?

An automated security testing framework that scans LLM applications for vulnerabilities:

  • โœ… 27 validated attack vectors across 6 categories
  • โœ… $0.50 average scan cost with built-in budget controls
  • โœ… 30-second results - Docker run, immediate feedback
  • โœ… CI/CD ready - Integrate into your deployment pipeline
  • โœ… Novel research - First tool to systematically test document-based instruction hijacking

๐Ÿš€ Quick Start

# Step 1: Install
pip install vektor

# Step 2: Zero-setup demo (no API key)
vektor demo

# Step 3: Real scan, $0 cost, always works
vektor scan --target vulnerable --output my-first-report.html

# Step 4: Open the report
# Windows: start my-first-report.html
# Mac:     open my-first-report.html
# Linux:   xdg-open my-first-report.html

# Step 5 (optional): Scan your own LLM app
export OPENAI_API_KEY=sk-your-key
vektor scan --target openai --budget 1.0

๐ŸŒ Scan Any AI API โ€” No SDK Needed

Point Vektor at any HTTP endpoint:

# Auto-detects OpenAI/Anthropic/custom shapes
vektor scan --url http://localhost:8000/chat

# With auth header
vektor scan --url https://my-app.com/api \
  --header "Authorization: Bearer YOUR_TOKEN"

# Custom request/response field names
vektor scan --url http://localhost:8000/predict \
  --request-field prompt --response-field answer

# Query-parameter mode (e.g. /api/parse?text=PAYLOAD)
vektor scan --url http://localhost:8000/api/parse \
  --param-field text

# Rate-limited API โ€” add delay between requests
vektor scan --url http://localhost:8000/chat \
  --request-delay 12.0

๐Ÿ’ก Why Vektor?

Feature Vektor Garak Promptfoo PyRIT
Primary Focus Actionable Security Vulnerability Scanning General Eval / Testing Red Teaming Framework
Setup Time < 30s ~10 mins ~5 mins ~30 mins
Scan Speed Fast (Targeted) Slow (Exhaustive) Fast Slow (Agentic)
Cost Control โœ… Built-in Budget โŒ โŒ โŒ
CI/CD Ready โœ… Native โš ๏ธ Heavy โœ… โš ๏ธ Complex
RAG/Doc Attacks โœ… Specialized โš ๏ธ Limited โœ… โœ…
Before After
Hire pentester ($5K) Run: vektor scan
Wait 2 weeks Get results in 1 minute
Get 50-page report Actionable JSON/HTML reports
Still don't know if fixes work Re-run to validate fixes

๐Ÿ”ฌ Attack Categories

1. Prompt Injection (6 attacks)

  • Direct instruction injection
  • System prompt override
  • Delimiter confusion
  • Role manipulation
  • Multi-turn context poisoning
  • Encoding-based bypass

2. Data Extraction (4 attacks)

  • Training data leak attempts
  • System prompt disclosure
  • Context window extraction
  • PII leakage testing

3. Instruction Hijacking (5 attacks) โ€” NOVEL

  • Simple document injection
  • DOCX hidden text injection
  • DOCX footnote injection
  • Markdown comment injection
  • Multi-document context poisoning

4. RAG Attacks (5 attacks)

  • Context poisoning via retrieved docs
  • RAG prompt leakage
  • Source fabrication / hallucination injection
  • Indirect injection via document store
  • Chunking boundary exploitation

5. Agent Attacks (4 attacks)

  • Tool call injection
  • Goal hijacking
  • Memory poisoning
  • Agent scope escape

6. Structured Output Injection (3 attacks)

  • JSON schema bypass
  • Output format injection
  • Type confusion attack

๐Ÿ“ฆ Installation

Docker

docker build -t vektor .
docker run -e OPENAI_API_KEY=$OPENAI_API_KEY vektor scan --target openai

pip

pip install vektor

From Source

git clone https://github.com/swapnilwankhede23/vektor.git
cd vektor
pip install -e .

๐Ÿ“– Usage

Basic Scan

vektor scan --target openai --budget 1.0

Quick Mode (High-confidence attacks only)

vektor scan --target openai --quick

CI/CD Integration

vektor scan --target openai --ci --output report.json

Specific Attacks

vektor scan --target openai --attacks direct_injection,system_override

Demo Mode (No API calls)

vektor demo

๐Ÿ“Š Sample Output

โ•”โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•—
โ•‘  V  E  K  T  O  R               โ•‘
โ•‘  AI Security Testing Framework   โ•‘
โ•šโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•

โ ‹ Testing attacks... โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ” 100% (15/15)

โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”ณโ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”ณโ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”“
โ”ƒ Attack                  โ”ƒ Severity  โ”ƒ Success    โ”ƒ
โ”กโ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ•‡โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ•‡โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”โ”ฉ
โ”‚ Direct Injection        โ”‚ HIGH      โ”‚ 83%        โ”‚
โ”‚ DOCX Hidden Text        โ”‚ CRITICAL  โ”‚ 67%        โ”‚
โ”‚ System Prompt Reveal    โ”‚ HIGH      โ”‚ 50%        โ”‚
โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜

โ•ญโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ•ฎ
โ”‚ Summary                         โ”‚
โ”œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ค
โ”‚ Risk Score: 72/100              โ”‚
โ”‚ Total Vulnerabilities: 3        โ”‚
โ”‚ Cost: $0.47                     โ”‚
โ”‚                                 โ”‚
โ”‚ HIGH RISK: Address all critical โ”‚
โ”‚ vulnerabilities before deploy   โ”‚
โ•ฐโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ•ฏ

โœ“ Report saved to: report.json

๐Ÿ” Novel Research: Instruction Hijacking

Vektor is the first tool to systematically test document-based instruction hijacking - a new class of vulnerabilities where attackers embed malicious instructions in document formats (DOCX, PDF, Markdown) that get processed by RAG systems.

Our research found:

  • 60% of RAG systems vulnerable to hidden text injection
  • 40% vulnerable to footnote/comment injection
  • Standard sanitization doesn't catch format-specific exploits

Read the full research paper โ†’

๐Ÿ› ๏ธ CI/CD Integration

GitHub Actions

name: LLM Security Scan
on: [push, pull_request]
jobs:
  security:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v3
      - uses: actions/setup-python@v4
        with:
          python-version: '3.11'
      - name: Install Vektor
        run: pip install vektor
      - name: Scan (no API key needed)
        run: vektor scan --target vulnerable --ci --output report.json
      - name: Upload Report
        uses: actions/upload-artifact@v3
        with:
          name: security-report
          path: report.json
      # Optional: scan your real LLM endpoint
      # - name: Scan real endpoint
      #   env:
      #     OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
      #   run: vektor scan --target openai --ci --output report.json

๐Ÿค Contributing

We welcome contributions! See CONTRIBUTING.md for guidelines.

Adding Custom Attacks

from vektor.attacks.base import BaseAttack, Vulnerability

class MyCustomAttack(BaseAttack):
    def __init__(self):
        super().__init__(name="my_attack", category="Custom")

    def execute(self, target):
        # Your attack logic
        pass

๐Ÿ“š Documentation

๐Ÿ—บ๏ธ Roadmap

v0.2 (Current)

  • โœ… 27 attack vectors across 6 categories
  • โœ… HTTP endpoint target (vektor scan --url http://localhost:8000/chat)
  • โœ… RAG pipeline targets (LangChain, LlamaIndex)
  • โœ… Agent targets (LangGraph, CrewAI, AutoGen)
  • โœ… Regression diff system for CI gating
  • โœ… Python scan() API
  • โœ… Docker deployment + CI/CD integration

v0.3 (Next)

  • โณ Web dashboard
  • โณ PDF document testing
  • โณ Multi-model comparison
  • โณ Compliance reporting (OWASP LLM Top 10 mapping)

๐Ÿ“„ License

MIT License - see LICENSE for details

๐Ÿ™ Acknowledgments

Built on research from:

  • Simon Willison (@simonw) - Prompt injection taxonomy
  • Greshake et al. - Indirect prompt injection
  • OWASP LLM Top 10 Project

โญ Support

If you find this useful, please star the repository and share with your network!

๐Ÿ“ง Contact


โš ๏ธ Disclaimer: This tool is for security testing purposes only. Use responsibly and only on systems you have permission to test.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

vektor_scan-0.2.0.tar.gz (78.3 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

vektor_scan-0.2.0-py3-none-any.whl (92.4 kB view details)

Uploaded Python 3

File details

Details for the file vektor_scan-0.2.0.tar.gz.

File metadata

  • Download URL: vektor_scan-0.2.0.tar.gz
  • Upload date:
  • Size: 78.3 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.7

File hashes

Hashes for vektor_scan-0.2.0.tar.gz
Algorithm Hash digest
SHA256 1dbd8803f3b158f91de2960e09258d07e93dacd3be5283f919c132df1a7a3d59
MD5 a724476b88438a51d9edcbf60b85db8f
BLAKE2b-256 486dc061da4352a23879b0d885160d5eec560d32ed951041d97b0b39f46dcf4a

See more details on using hashes here.

File details

Details for the file vektor_scan-0.2.0-py3-none-any.whl.

File metadata

  • Download URL: vektor_scan-0.2.0-py3-none-any.whl
  • Upload date:
  • Size: 92.4 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.7

File hashes

Hashes for vektor_scan-0.2.0-py3-none-any.whl
Algorithm Hash digest
SHA256 d30fb1100b45649292fb429e4d0df813545e6135d26c0bbafc74917fb2501255
MD5 bf0569c2c365093c726a3412378925dc
BLAKE2b-256 fab1e52a2cc8f9d66a6727d2371aed43c6775743a72101e727d91594f98f60ce

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page