AIDR Bastion

A comprehensive GenAI protection system designed to safeguard against malicious prompts, injection attacks, and harmful content. The system incorporates multiple detection engines that operate sequentially to analyze and classify user inputs before reaching GenAI applications.

• The system supports Roota and Sigma rules, enabling the application of detection logic from multiple sources such as SigmaHQ (around 1,200 compatible free community Sigma rules available at release), SOC Prime (with up to 3,000 additional compatible rules), and other third-party repositories. Sigma rules can be applied to detect use cases where malware leverages a local LLM to generate malicious code for execution.
• SOC Prime Uncoder AI integration further extends functionality by translating Sigma rules into Semgrep format, providing standardized and reusable detection pipelines (requires a free account).
• Roota rules power the regex-based pipeline.
• The architecture supports rule extensibility, seamlessly integrating organization-specific signatures and external detection content.
• The system can also function as a local logging sensor, recording user and agent prompts and enabling diagnostics, incident discovery, and cyber attack investigation.
• Detection logic aligns with industry frameworks such as MITRE ATLAS and OWASP Top 10 for LLMs, ensuring standardized coverage against adversarial techniques.
• Actions include allow, block, or notify, depending on rule matches and policy configuration.

This layered detection approach delivers defense-in-depth against evolving adversarial prompt engineering and other AI-focused attack vectors. Inspired by LlamaFirewall.

🚀 Features

• Multi-Pipeline Detection: Regex patterns, ML models, vector-based similarity detection, and LLM-based analysis
• Flexible Configuration: Dynamic Pipeline configuration via JSON
• Real-time Analysis: Fast async processing with configurable thresholds
• Client Managers: Flexible client management (Elasticsearch, OpenSearch, Qdrant)
• RESTful API: Easy integration with existing applications
• Extensible Architecture: Simple plugin system for custom Pipelines

🏗️ Architecture

┌────────────────────────────────┐
│   FastAPI Endpoint             │
│   (POST /api/v1/flow/run)      │
└──────────────┬─────────────────┘
               │
               ▼
      ┌─────────────────────┐
      │  Pipeline Manager   │
      └─────────┬───────────┘
                │
                ▼
      ┌──────────────────────────────┐
      │          Pipelines           │
      │ ┌──────────────────────────┐ │
      │ │  Rule Pipeline           │ │
      │ ├──────────────────────────┤ │
      │ │  Similarity Pipeline     │ │
      │ │  (Similarity Manager)    │ │
      │ ├──────────────────────────┤ │
      │ │  Code Analysis Pipeline  │ │
      │ ├──────────────────────────┤ │
      │ │  ML Pipeline             │ │
      │ ├──────────────────────────┤ │
      │ │  LLM Pipeline            │ │
      │ │  (LLM Manager)           │ │
      │ └──────────────────────────┘ │
      └──────────────────────────────┘

📋 Table of Contents

• Installation
• Configuration
• Usage
• API Reference
• Pipelines
• Managers
• Rule Management and Customization
• Adding Custom Pipelines
• Development
• License
• Built With
• TO-DO List

🛠️ Installation

Prerequisites

• Python 3.12+
• OpenSearch, Elasticsearch, or Qdrant (via Similarity Manager)
• LLM API key (optional, for LLM Pipeline):
  • OpenAI API key for GPT models, or
  • Anthropic API key for Claude models, or
  • Azure OpenAI credentials for enterprise deployments, or
  • Ollama running locally for privacy-focused local LLM models

Quick Start

1. Clone the repository

   git clone https://github.com/0xAIDR/AIDR-Bastion.git
   cd AIDR-Bastion
   

2. Create virtual environment

   python -m venv venv
   source venv/bin/activate  # On Windows: venv\Scripts\activate
   

3. Install dependencies

   pip install -r requirements.txt
   

4. Set up environment variables

   cp env.example .env
   # Edit .env with your configuration
   

5. Run the application

   python server.py
   

By default, the API will be available at http://localhost:8000. You can customize the host and port using the HOST and PORT environment variables.

⚙️ Configuration

Environment Variables (.env)

# FastAPI configuration
HOST=0.0.0.0
PORT=8000


# ML Pipeline. 
# Path to the model
ML_MODEL_PATH=

# LLM Pipeline
LLM_DEFAULT_CLIENT=openai  # openai, anthropic, azure, or ollama

# OpenAI Configuration
OPENAI_API_KEY=
OPENAI_MODEL=gpt-4
OPENAI_BASE_URL=https://api.openai.com/v1

# Anthropic Configuration
ANTHROPIC_API_KEY=
ANTHROPIC_MODEL=claude-sonnet-4-5-20250929
ANTHROPIC_BASE_URL=https://api.anthropic.com

# Azure OpenAI Configuration
AZURE_OPENAI_ENDPOINT=https://your-resource.openai.azure.com/
AZURE_OPENAI_API_KEY=
AZURE_OPENAI_DEPLOYMENT=gpt-4
AZURE_OPENAI_API_VERSION=2024-02-15-preview

# Ollama Configuration (for local LLM models)
OLLAMA_BASE_URL=http://localhost:11434
OLLAMA_MODEL=llama3

# LLM Common Configuration (applies to all LLM providers)
LLM_TEMPERATURE=0.1  # Temperature for LLM responses (0.0-2.0, lower = more focused)
LLM_MAX_TOKENS=1000  # Maximum tokens for LLM responses

# Similarity Pipeline
# similarity-prompt-index by default
SIMILARITY_PROMPT_INDEX=

SIMILARITY_NOTIFY_THRESHOLD=0.7
SIMILARITY_BLOCK_THRESHOLD=0.87

# Manager configuration
SIMILARITY_DEFAULT_CLIENT=opensearch  # opensearch, elasticsearch, or qdrant
LLM_DEFAULT_CLIENT=openai  # openai, anthropic, azure, or ollama

# OpenSearch configuration
OS__HOST=
OS__PORT=
OS__SCHEME=
OS__USER=
OS__PASSWORD=

# Elasticsearch configuration (alternative to OpenSearch)
ES__HOST=
ES__PORT=
ES__SCHEME=
ES__USER=
ES__PASSWORD=

# Qdrant configuration (alternative to OpenSearch/Elasticsearch)
QDRANT__HOST=localhost
QDRANT__PORT=6333
QDRANT__GRPC_PORT=6334
QDRANT__API_KEY=
QDRANT__PREFER_GRPC=false
QDRANT__TIMEOUT=30

# Kafka configuration (for event logging)
KAFKA__BOOTSTRAP_SERVERS=localhost:9092
KAFKA__TOPIC=aidr-events
KAFKA__SECURITY_PROTOCOL=PLAINTEXT
KAFKA__SASL_MECHANISM=
KAFKA__SASL_USERNAME=
KAFKA__SASL_PASSWORD=

# requires for creating embedding in pipelines: Similarity Pipeline and ML Pipeline
EMBEDDINGS_MODEL=

## Kafka configuration
# KAFKA__BOOTSTRAP_SERVERS=
# KAFKA__TOPIC=
# KAFKA__SECURITY_PROTOCOL=PLAINTEXT
# KAFKA__SASL_MECHANISM=
# KAFKA__SASL_USERNAME=
# KAFKA__SASL_PASSWORD=
# KAFKA__SAVE_PROMPT=true 

Pipeline Configuration (config.json)

The config.json file controls which Pipelines will be run for each flow. Default config.json configuraton:

[
    {
        "pipeline_flow": "full_scan",
        "pipelines": [
            "similarity",
            "rule",
            "openai",
            "ml",
            "code_analysis"
        ]
    },
    {
        "pipeline_flow": "code_audit",
        "pipelines": [
            "code_analysis"
        ]
    },
    {
        "pipeline_flow": "model_audit",
        "pipelines": [
            "ml",
            "openai"
        ]
    },
    {
        "pipeline_flow": "base_audit",
        "pipelines": [
            "rule",
            "similarity"
        ]
    }
]

Configuration Impact

• Flow names: Can be any custom name (e.g., base, code, security, content). The name must match what you pass in the API request's pipeline_flow parameter
• Pipeline names: Must match the Pipeline names defined in PipelineNames enum
• Order matters: Pipelines run in the order specified in the array
• Example flows:
  • base flow: Pipelines general text prompts for harmful content
  • code flow: Pipelines code snippets for security vulnerabilities
  • custom_flow: You can create any custom flow name for specific use cases

🚀 Usage

Basic API Usage

import requests

# Run pipeline analysis on a text prompt
response = requests.post("http://localhost:8000/api/v1/flow/run", json={
    "prompt": "Your text to analyze here",
    "pipeline_flow": "base"  # Must match a flow_name from config.json
})

result = response.json()
print(f"Status: {result['status']}")  # allow, block, or notify
print(f"Triggered rules: {result['result']}")

# Get available flows and pipelines
flows_response = requests.get("http://localhost:8000/api/v1/flow/list")
flows = flows_response.json()
print(f"Available flows: {[flow['flow_name'] for flow in flows['flows']]}")

# Get available managers and their clients
managers_response = requests.get("http://localhost:8000/api/v1/manager/list")
managers = managers_response.json()
print(f"Available managers: {[manager['name'] for manager in managers['managers']]}")

# Get information about a specific manager
similarity_manager = requests.get("http://localhost:8000/api/v1/manager/similarity")
manager_info = similarity_manager.json()
print(f"Similarity Manager clients: {manager_info['clients']}")

# Switch active client for a manager
switch_response = requests.post("http://localhost:8000/api/v1/manager/switch_active_client", json={
    "manager_id": "similarity",
    "client_id": "elasticsearch"
})
switch_result = switch_response.json()
print(f"Client switched: {switch_result['status']}")

Python SDK Usage

from app.main import bastion_app

# Direct usage
result = await bastion_app.run("Your prompt", "default")
print(f"Status: {result.status}")
for pipeline in result.pipelines:
    print(f"Pipeline: {pipeline._identifier}, Status: {pipeline.status}")

Integration with Existing Applications

You can integrate project for your existing LLM application:

1. Send requests:

import requests

def check_prompt_safety(prompt: str):
    response = requests.post(
        "http://localhost:8000/api/v1/flow/run",
        json={
            "prompt": prompt,
            "pipeline_flow": "security_audit"
        }
    )
    result = response.json()
    
    if result["status"] == "BLOCK":
        return False, "Prompt blocked"
    elif result["status"] == "NOTIFY":
        return True, "Prompt flagged but allowed"
    else:
        return True, "Prompt safe"

2. Configure your application to check all user inputs
3. Set up proper error handling and fallbacks
4. Manage managers and clients dynamically:

import requests

def get_available_clients():
    """Get list of available clients for each manager"""
    response = requests.get("http://localhost:8000/api/v1/manager/list")
    managers = response.json()
    
    for manager in managers['managers']:
        print(f"{manager['name']}: {manager['clients']}")
        print(f"Enabled: {manager['enabled']}")

def switch_to_elasticsearch():
    """Switch Similarity Manager to use Elasticsearch"""
    response = requests.post("http://localhost:8000/api/v1/manager/switch_active_client", json={
        "manager_id": "similarity",
        "client_id": "elasticsearch"
    })
    result = response.json()
    return result['status']

Project Configuration

The project can be configured through environment variables:

• HOST: Server host (default: 0.0.0.0)
• PORT: Server port (default: 8000)
• CORS_ORIGINS: Allowed origins for CORS
• EMBEDDINGS_MODEL: Hugging Face model for embeddings
• SIMILARITY_NOTIFY_THRESHOLD: Threshold for notifications
• SIMILARITY_BLOCK_THRESHOLD: Threshold for blocking

All required environments you can find in env.example

📚 API Reference

POST /api/v1/flow/run

Runs pipelines to analyze the input prompt.

Request Body:

{
    "prompt": "string",
    "pipeline_flow": "string"  // Must match a flow_name from config.json
}

Response:

{
    "status": "allow | block | notify",
    "result": [
        {
            "status": "allow | block | notify",
            "name": "string",
            "triggered_rules": [
                {
                    "id": "string",
                    "name": "string",
                    "details": "string",
                    "body": "string",
                    "action": "notify" | "block",
                    "severity": "string",
                    "cwe_id": "string"
                }
            ]
        }
    ]
}

GET /api/v1/flow/list

Get a list of all available flows and their pipelines.

Response:

{
    "flows": [
        {
            "flow_name": "string",
            "pipelines": [
                {
                    "name": "string",
                    "enabled": "boolean"
                }
            ]
        }
    ]
}

GET /api/v1/manager/list

Get a list of all available managers and their clients.

Response:

{
    "managers": [
        {
            "id": "string",
            "name": "string",
            "enabled": "boolean",
            "clients": ["string"]
        }
    ]
}

GET /api/v1/manager/{manager_id}

Get information about a specific manager.

Parameters:

• manager_id (string): ID of the manager (e.g., "similarity", "llm")

Response:

{
    "id": "string",
    "name": "string",
    "enabled": "boolean",
    "clients": ["string"]
}

POST /api/v1/manager/switch_active_client

Switch the active client for a specific manager.

Request Body:

{
    "manager_id": "string",
    "client_id": "string"
}

Response:

{
    "client_id": "string",
    "status": "boolean"
}

🔍 Pipelines

1. Rule Pipeline (rule)

• Purpose: Pattern-based detection using regular expressions
• Rules: YAML files in app/pipelines/rule_pipeline/rules/
• Categories:
  • Injection: SQL, command execution, path traversal
  • Obfuscation: Character obfuscation, encoding, Unicode homoglyphs
  • Override: Role play, filter disabling, context splicing
  • Leakage: Direct prompt requests, forced repetition
  • PII: Email, phone, credit cards, passwords, API keys, UUIDs, IBAN
  • Semantic: Emotional manipulation, authority fallacy, multilingual attacks
  • DoS: Character/word repetition, regex DoS
• Best for: Known attack patterns and simple text analysis

2. Similarity Pipeline (similarity)

• Purpose: Vector-based similarity detection against known harmful prompts
• Backend: Uses Similarity Manager for vector search
• Required: OpenSearch, Elasticsearch, or Qdrant configuration via Similarity Manager
• Configuration: SIMILARITY_NOTIFY_THRESHOLD, SIMILARITY_BLOCK_THRESHOLD, SIMILARITY_DEFAULT_CLIENT
• Best for: Detecting variations of known attacks

3. Code Analysis Pipeline (code_analysis)

• Purpose: Static code analysis using Semgrep
• Languages: Python, JavaScript, Java, C++, and more
• Rules: Security-focused patterns
• Best for: Code injection and vulnerability detection

4. ML Pipeline (ml)

• Purpose: Machine learning-based classification
• Configuration: Requires ML_PIPELINE_PATH
• Model: Custom-trained model for prompt classification
• Best for: General malicious content detection
• Required: Configured environment EMBEDDINGS_MODEL

5. LLM Pipeline (llm)

• Purpose: AI-powered analysis using advanced language models
• Backend: Uses LLM Manager for text analysis
• Supported Providers:
  • OpenAI - GPT-4, GPT-4 Turbo models
  • Anthropic - Claude 3.5 Sonnet, Claude 3 Opus models
  • Azure OpenAI - Enterprise GPT models via Microsoft Azure infrastructure
  • Ollama - Local LLM models (Llama 3, Mistral, Gemma, etc.) for privacy-focused deployments
• Configuration:
  • LLM_DEFAULT_CLIENT - Choose provider (openai, anthropic, azure, ollama)
  • LLM_TEMPERATURE - Control response randomness (0.0-2.0, default: 0.1)
  • LLM_MAX_TOKENS - Maximum response length (default: 1000)
  • Provider-specific API keys and settings
• Features: JSON response format, configurable models, intelligent decision-making, multi-provider support
• Response Format: Returns structured JSON with status (block/notify/allow) and reasoning
• Best for: Complex reasoning, context-aware analysis, and nuanced content moderation

👥 Managers

The system uses managers to control different types of clients and provide a unified interface for pipelines.

Available managers:

• Similarity Manager - manages search clients for vector search
• LLM Manager - manages LLM clients for text analysis

Similarity Manager

Manages search systems for vector search of similar content. Automatically selects available client based on configuration.

Available clients:

• OpenSearch Client - Full-text search with KNN plugin for vector similarity
• Elasticsearch Client - Search engine with dense_vector support for cosine similarity
• Qdrant Client - Specialized vector database optimized for high-performance similarity search

Use the SIMILARITY_DEFAULT_CLIENT environment variable to set the default client. The endpoint /api/v1/manager/switch_active_client also allows this.

Qdrant advantages:

• 2-3x faster vector search with native HNSW algorithm
• Lower memory footprint and better resource efficiency
• Simpler API without complex KNN queries
• Built-in score threshold filtering
• Native support for gRPC protocol
• Ideal for high-throughput production environments

Clients in development:

• Planned support for other vector databases (PostgreSQL)
• You can contribute clients

LLM Manager

Manages LLM providers for text analysis and classification. Provides unified interface for multiple AI providers with dynamic client switching.

Available clients:

• OpenAI Client - GPT-4, GPT-4 Turbo models support
• Anthropic Client - Claude 3.5 Sonnet, Claude 3 Opus models support
• Azure OpenAI Client - Enterprise GPT models via Microsoft Azure infrastructure
• Ollama Client - Local LLM models (Llama 3, Mistral, Gemma, etc.) for privacy-focused deployments

Use the LLM_DEFAULT_CLIENT environment variable to set the default client. The endpoint /api/v1/manager/switch_active_client allows dynamic switching between providers.

Example switching between providers:

# Switch to Anthropic Claude
requests.post("http://localhost:8000/api/v1/manager/switch_active_client", json={
    "manager_id": "llm",
    "client_id": "anthropic"
})

# Switch to OpenAI GPT
requests.post("http://localhost:8000/api/v1/manager/switch_active_client", json={
    "manager_id": "llm",
    "client_id": "openai"
})

# Switch to Azure OpenAI
requests.post("http://localhost:8000/api/v1/manager/switch_active_client", json={
    "manager_id": "llm",
    "client_id": "azure"
})

# Switch to Ollama (local)
requests.post("http://localhost:8000/api/v1/manager/switch_active_client", json={
    "manager_id": "llm",
    "client_id": "ollama"
})

Clients in development:

• Planned support for other LLM providers (Google Gemini, Groq, Mistral)
• You can contribute clients

📋 Rule Management and Customization

YAML Rules for Rule Pipeline

The Rule Pipeline defines detection patterns using Roota rules files. Roota is a public-domain language for collective cyber defense that combines native queries from SIEM, EDR, XDR, or Data Lake with standardized metadata and threat intelligence to enable automated translation into other languages.

Each rule file follows a specific structure:

name: "Rule Name"
description: "Description of what this rule detects"
severity: "high|medium|low"
category: "injection|obfuscation|override|leakage|pii|semantic|dos"
patterns:
  language: llm-regex-pattern
  pattern: 
   - "pattern"
   - "another_pattern"
action: "block|notify|allow"

Rule Categories:

• Injection: SQL injection, command execution, path traversal, script injection
• Obfuscation: Character obfuscation, encoding tricks, Unicode homoglyphs
• Override: Role play attacks, filter disabling, context splicing
• Leakage: Direct prompt requests, forced repetition attacks
• PII: Personal identifiable information detection (emails, phones, credit cards, etc.)
• Semantic: Emotional manipulation, authority fallacy, multilingual attacks
• DoS: Denial of service patterns (character repetition, regex DoS)

Semgrep Rules for Code Analysis Pipeline

The code pipeline uses Semgrep rules for static code analysis. Rules are located in app/pipelines/semgrep_pipeline/rules/.

Rule Structure:

rules:
  - id: rule-id
    message: "Security issue description"
    languages: [python, javascript, java]
    severity: ERROR
    patterns:
      - pattern: |
          $PATTERN
    fix: |
      $FIX

Managing Rules

Using Roota for Rule Creation

Roota is a public-domain language for collective cyber defense that provides:

• YAML-based format that's easy to write and human-readable
• Multi-language support for SIEM, EDR, XDR, and Data Lake queries
• MITRE ATT&CK mapping for threat intelligence integration
• Threat actor timeline support for coordinated defense
• Correlation support for more robust detection logic
• OCSF and Sigma compatibility for maximum compatibility

Roota Rule Example:

name: 'INJ-001: SQL Keywords'
details: Detects common SQL manipulation keywords. Designed to be a high-confidence signal. https://tdm.socprime.com/
author: SOC Prime Team
severity: critical
date: 2025-08-08
logsource:
  product: llm
  service: firewall
  module: regex
detection:
  language: llm-regex-pattern
  pattern:
    - '(?i)\b(?:SELECT\s+(?:(?!\bFROM\b)[^,;]+,)+(?:(?!\bFROM\b)[^,;]+)\s+FROM|INSERT\s+INTO|UPDATE\s+[\w\.]+\s+SET|DELETE\s+FROM|DROP\s+(?:TABLE|DATABASE)|ALTER\s+TABLE|CREATE\s+TABLE|TRUNCATE\s+TABLE)\b'
references:
  - https://genai.owasp.org/llmrisk/llm01-prompt-injection/
  - https://owasp.org/Top10/A03_2021-Injection/
license: DRL 1.1
uuid: f1a2b3c4-d5e6-4f7a-8b8c-9d0e1f2a3b4c
response: block

Using Uncoder AI for Semgrep Rules

Uncoder AI is a powerful tool for converting Sigma and Roota rules to various formats including Semgrep:

1. Visit Uncoder AI
2. Register account for free
3. Select Roota/Sigma to Semgrep conversion
4. Paste your Roota or Sigma rule
5. Generate Semgrep YAML rule
6. Save the generated rule in app/pipelines/semgrep_pipeline/rules/

Example Sigma Rule:

title: Suspicious PowerShell Command
description: Detects suspicious PowerShell commands
logsource:
  category: process_creation
  product: windows
detection:
  selection:
    - CommandLine: '*powershell*'
    - CommandLine: '*Invoke-Expression*'
  condition: selection

Using SOC Prime for Advanced Rules

SOC Prime provides comprehensive threat detection rules including Roota and Sigma formats:

1. Visit SOC Prime
2. Browse the Rules Library (Sigma, Roota, and other formats)
3. Filter by technology and threat type
4. Download or convert rules using Uncoder AI
5. Adapt rules for your specific use case

Custom Rule Development

Creating Custom regex Rules

1. Identify the attack pattern
2. Create YAML file in appropriate category folder
3. Define patterns with clear descriptions
4. Test thoroughly with various inputs
5. Set appropriate severity and action

Example Custom Rule:

name: Custom Injection Pattern
details: Detects custom injection attempts
author: your name
severity: high
date: 2025-08-08
logsource:
    product: llm
    category: injection
    module: regex
patterns:
    language: llm-regex-pattern
    pattern: 
        - "(?i)(union|select|insert|delete|update|drop).*from"
        - "(?i)(exec|system|eval|shell_exec)"
references:
    - https://one_example
    - https://two_example
license: DRL 1.1
uuid: f1a2b3c4-d5e6-4f7a-8b8c-9d0e1f2a3b4c
action: block

Creating Custom Semgrep Rules

1. Identify code vulnerability pattern
2. Write Semgrep pattern using their syntax
3. Test with sample code
4. Add appropriate metadata
5. Place in rules directory

Example Custom Semgrep Rule:

rules:
  - id: custom-sql-injection
    message: "Potential SQL injection vulnerability"
    languages: [python]
    severity: ERROR
    patterns:
      - pattern: |
          $QUERY = "SELECT * FROM $TABLE WHERE id = " + $USER_INPUT
    fix: |
      $QUERY = "SELECT * FROM $TABLE WHERE id = %s"
      # Use parameterized queries

⚙️ Client Selection

Priority Client Selection

• Similarity Manager: Uses SIMILARITY_DEFAULT_CLIENT to choose between OpenSearch, Elasticsearch, or Qdrant
• LLM Manager: Uses LLM_DEFAULT_CLIENT to choose LLM provider
• Dynamic switching: Can change active client via API endpoint /api/v1/manager/switch_active_client

Priority Configuration

# Priority for Similarity Manager
SIMILARITY_DEFAULT_CLIENT=opensearch  # opensearch, elasticsearch, or qdrant

# Priority for LLM Manager
LLM_DEFAULT_CLIENT=openai  # openai, anthropic, azure, or ollama

⚙️ Enabling Disabled Pipeline

Some pipelines are disabled by default. To enable them:

Enable LLM Pipeline with OpenAI

1. Configure your OpenAI API key in the .env configuration file:

   LLM_DEFAULT_CLIENT=openai
   OPENAI_API_KEY=your-api-key
   OPENAI_MODEL=gpt-4
   OPENAI_BASE_URL=https://api.openai.com/v1
   
   # Optional: LLM Common Configuration (applies to all providers)
   LLM_TEMPERATURE=0.1  # Lower = more focused, higher = more creative
   LLM_MAX_TOKENS=1000  # Maximum response length
   

2. Add to your flow in config.json:

   {
       "flow_name": "base",
       "pipelines": ["similarity", "rule", "llm"]
   }
   

Enable LLM Pipeline with Anthropic Claude

1. Configure your Anthropic API key in the .env configuration file:

   LLM_DEFAULT_CLIENT=anthropic
   ANTHROPIC_API_KEY=your-api-key
   ANTHROPIC_MODEL=claude-sonnet-4-5-20250929
   ANTHROPIC_BASE_URL=https://api.anthropic.com
   
   # Optional: LLM Common Configuration (applies to all providers)
   LLM_TEMPERATURE=0.1  # Lower = more focused, higher = more creative
   LLM_MAX_TOKENS=1000  # Maximum response length
   

2. Add to your flow in config.json:

   {
       "flow_name": "base",
       "pipelines": ["similarity", "rule", "llm"]
   }
   

Enable LLM Pipeline with Azure OpenAI

1. Configure your Azure OpenAI credentials in the .env configuration file:

   LLM_DEFAULT_CLIENT=azure
   AZURE_OPENAI_ENDPOINT=https://your-resource.openai.azure.com/
   AZURE_OPENAI_API_KEY=your-api-key
   AZURE_OPENAI_DEPLOYMENT=gpt-4
   AZURE_OPENAI_API_VERSION=2024-02-15-preview
   
   # Optional: LLM Common Configuration (applies to all providers)
   LLM_TEMPERATURE=0.1  # Lower = more focused, higher = more creative
   LLM_MAX_TOKENS=1000  # Maximum response length
   

2. Add to your flow in config.json:

   {
       "flow_name": "base",
       "pipelines": ["similarity", "rule", "llm"]
   }
   

Enable LLM Pipeline with Ollama (Local Models)

1. Install and start Ollama:

   # Install Ollama from https://ollama.ai
   # Pull a model (e.g., Llama 3)
   ollama pull llama3
   
   # Ollama automatically starts on http://localhost:11434
   

2. Configure Ollama in the .env configuration file:

   LLM_DEFAULT_CLIENT=ollama
   OLLAMA_BASE_URL=http://localhost:11434
   OLLAMA_MODEL=llama3
   
   # Optional: LLM Common Configuration (applies to all providers)
   LLM_TEMPERATURE=0.1  # Lower = more focused, higher = more creative
   LLM_MAX_TOKENS=1000  # Maximum response length
   

3. Add to your flow in config.json:

   {
       "flow_name": "base",
       "pipelines": ["similarity", "rule", "llm"]
   }
   

4. Switch between providers dynamically via API:

   # Switch to Anthropic
   requests.post("http://localhost:8000/api/v1/manager/switch_active_client",
       json={"manager_id": "llm", "client_id": "anthropic"})
   
   # Switch to OpenAI
   requests.post("http://localhost:8000/api/v1/manager/switch_active_client",
       json={"manager_id": "llm", "client_id": "openai"})
   
   # Switch to Azure OpenAI
   requests.post("http://localhost:8000/api/v1/manager/switch_active_client",
       json={"manager_id": "llm", "client_id": "azure"})
   
   # Switch to Ollama
   requests.post("http://localhost:8000/api/v1/manager/switch_active_client",
       json={"manager_id": "llm", "client_id": "ollama"})
   

Enable ML Pipeline

1. Configure the model path in the .env configuration file:

   ML_MODEL_PATH=/path/to/your/model
   

2. Add to your flow in config.json:

   {
       "flow_name": "base",
       "pipelines": ["similarity", "rule", "ml"]
   }
   

🔧 Adding Custom Pipelines

Step 1: Create Pipeline Class

# app/pipelines/my_pipeline/pipeline.py
from app.core.pipeline import BasePipeline
from app.core.enums import PipelineNames, ActionStatus
from app.models.pipeline import PipelineResult, TriggeredRuleData

class MyCustomPipeline(BasePipeline):
    name = PipelineNames.my_pipeline
    enabled = True

    async def run(self, prompt: str, **kwargs) -> PipelineResult:
        # Your analyzing logic here
        triggered_rules = []
        
        # Example: Check for specific patterns
        if "malicious_pattern" in prompt.lower():
            triggered_rules.append(TriggeredRuleData(
                id="my_rule_1",
                name="Malicious Pattern Detected",
                details="Found potentially malicious content",
                body=prompt,
                action=RuleAction.BLOCK
            ))

        status = ActionStatus.BLOCK if triggered_rules else ActionStatus.ALLOW
        return PipelineResult(
            name=self._identifier,
            status=status,
            triggered_rules=triggered_rules
        )

Step 2: Register Pipeline

# app/pipelines/__init__.py
from app.pipelines.my_pipeline.pipeline import MyCustomPipeline

__PIPELINES__ = [
    # ... existing pipelines
    MyCustomPipeline(),
]

PIPELINES_MAP = {
    pipeline._identifier: pipeline for pipeline in __PIPELINES__
    if pipeline.enabled
}

Step 3: Add to Configuration

[
    {
        "flow_name": "base",
        "pipelines": [
            "personal_info",
            "similarity",
            "rule",
            "my_pipeline"
        ]
    }
]

Step 4: Add Pipeline Name to Enum

# app/core/enums.py
class PipelineNames(str, Enum):
    # ... existing names
    my_pipeline = "my_pipeline"

🧪 Development

Setting up OpenSearch

1. Install OpenSearch

   docker run -p 9200:9200 -p 9600:9600 -e "discovery.type=single-node" opensearchproject/opensearch:latest
   

2. Create similarity index

   python app/scripts/similarity/index_script.py
   

   This will create the similarity-prompt-index index in OpenSearch. You can customize the index name by setting the SIMILARITY_PROMPT_INDEX environment variable.

Setting up Elasticsearch

1. Install Elasticsearch

   # Alternative to OpenSearch
   docker run -p 9200:9200 -e "discovery.type=single-node" elasticsearch:latest
   

2. Create similarity index

   python app/scripts/similarity/index_script.py
   

   This will create the similarity-prompt-index index in Elasticsearch. You can customize the index name by setting the SIMILARITY_PROMPT_INDEX environment variable.

Setting up Qdrant

1. Install Qdrant

   # Recommended: Fastest and most efficient option
   docker run -p 6333:6333 -p 6334:6334 qdrant/qdrant
   

2. Configure environment variables

   # Add to your .env file
   SIMILARITY_DEFAULT_CLIENT=qdrant
   QDRANT__HOST=localhost
   QDRANT__PORT=6333
   QDRANT__GRPC_PORT=6334
   # Optional: for Qdrant Cloud
   # QDRANT__API_KEY=your-api-key
   # QDRANT__PREFER_GRPC=true
   

3. Create similarity collection

   python scripts/similarity/index_script.py
   

   This will create the similarity-prompt-index collection in Qdrant with optimized vector search configuration (HNSW algorithm, cosine similarity).

Why choose Qdrant:

• Performance: 2-3x faster than OpenSearch/Elasticsearch for vector operations
• Efficiency: Lower memory usage with optimized vector storage
• Simplicity: Cleaner API, no complex query DSL required
• Features: Built-in payload filtering, quantization, and snapshots
• Scale: Production-ready with horizontal scaling and cloud options

Setting up Kafka for Event Logging

AIDR Bastion supports Kafka integration for logging BLOCK and NOTIFY events, enabling scalable event streaming and real-time monitoring.

Quick Start with Docker Compose

Configure environment variables Minimal required environments

# Add to your .env file
KAFKA__BOOTSTRAP_SERVERS=localhost:9092
KAFKA__TOPIC=aidr-events
KAFKA__SECURITY_PROTOCOL=PLAINTEXT

Full Kafka environment variables

## Kafka configuration
# KAFKA__BOOTSTRAP_SERVERS=
# KAFKA__TOPIC=
# KAFKA__SECURITY_PROTOCOL=PLAINTEXT
# KAFKA__SASL_MECHANISM=
# KAFKA__SASL_USERNAME=
# KAFKA__SASL_PASSWORD=
# KAFKA__SAVE_PROMPT=true 

The environment variable KAFKA__SAVE_PROMPT is optional. It controls whether the input prompt data should be saved to Kafka or not.

Event Logging Features

• BLOCK Events: Logged when prompts are blocked by detection rules
• NOTIFY Events: Logged when prompts trigger notifications but are allowed
• Structured JSON: Events include prompt content, detection results, and metadata
• Real-time Streaming: Events are sent immediately to Kafka topics

Event Schema

{
	"status": "block",
	"pipelines": [
		{
			"status": "block",
			"name": "Pipeline Name",
			"triggered_rules": [
				{
					"details": "",
					"action": "block",
					"id": "a12d86d8-d96a-41fa-9e9a-18231539cfde",
					"name": "Instruction Overriding",
					"severity": null,
					"cwe_id": null
				}
			]
		}
	],
	"service": "AIDR Bastion",
	"version": "1.0.0",
	"timestamp": "2025-09-24T14:39:50.351466",
	"task_id": 1 // unique identifier passed through endpoint
}

Advanced Kafka Configuration

For production environments, configure additional security settings:

# SASL Authentication
KAFKA__SECURITY_PROTOCOL=SASL_SSL
KAFKA__SASL_MECHANISM=PLAIN
KAFKA__SASL_USERNAME=your_username
KAFKA__SASL_PASSWORD=your_password

# SSL Configuration (if required)
KAFKA__SSL_CA_LOCATION=/path/to/ca-cert
KAFKA__SSL_CERTIFICATE_LOCATION=/path/to/client-cert
KAFKA__SSL_KEY_LOCATION=/path/to/client-key

Rule Management

• test_rules.py: Comprehensive rule testing and validation
• generate_rules.py: Interactive rule creation and conversion
• Roota: Public-domain language for collective cyber defense
• Uncoder AI: Convert Roota/Sigma rules to Semgrep format
• SOC Prime: Access comprehensive threat detection rules

📄 License

This project is licensed under the GNU Lesser General Public License v3.0 (LGPL-3.0) - see the LICENSE file for details.

For more information about LGPL, visit: https://www.gnu.org/licenses/lgpl-3.0.html

🛠️ Built With

This project is built using the following powerful open-source libraries and frameworks:

• Roota - Public-domain language for collective cyber defense
• Uncoder AI - Convert Roota/Sigma rules to Semgrep format
• FastAPI - Modern, fast web framework for building APIs with Python 3.7+ based on standard Python type hints
• OpenSearch - Open source search and analytics suite for log analytics, application search, and more
• Elasticsearch - Open search and analytics platform for various data types
• OpenAI - AI research company providing powerful language models (GPT-4, GPT-4 Turbo) for intelligent content analysis
• Anthropic - AI safety company providing Claude models (Claude 3.5 Sonnet, Claude 3 Opus) for advanced reasoning and content moderation
• Azure OpenAI - Enterprise-grade OpenAI models via Microsoft Azure infrastructure with enhanced security and compliance
• Ollama - Run large language models locally with privacy-focused deployments (Llama 3, Mistral, Gemma, and more)
• Qdrant - High-performance vector search engine optimized for similarity search and filtering
• Semgrep - Static analysis tool for finding bugs and security issues in code
• Sentence Transformers - Python framework for state-of-the-art sentence, text and image embeddings
• Pydantic - Data validation and settings management using Python type annotations
• Uvicorn - Lightning-fast ASGI server implementation
• PyYAML - YAML parser and emitter for Python
• NLTK - Natural Language Toolkit for text processing

🛠️ TO-DO List

• Integrate API with SOC Prime for automatic rule synchronization and uploads
• Add local database storage for rules and events
• ✅ Kafka support for scalable event streaming - COMPLETED
• Develop an admin panel for managing events and detection rules
• Explore integration with NOVA Rules to extend rule sources
• Add YARA-L support