socprime/AIDR-Bastion
GitHub: socprime/AIDR-Bastion
一款面向大模型应用的多层安全防护网关，通过规则引擎、向量检索、ML 分类和 LLM 分析串联检测恶意 prompt 与注入攻击。
Stars: 98 | Forks: 7
# AIDR Bastion
[](VERSION)
[](https://www.python.org/downloads/release/python-3120/)
[](https://fastapi.tiangolo.com/)
[](LICENSE)
一个全面的 GenAI 保护系统，旨在防御恶意 prompt、注入攻击和有害内容。该系统包含多个按顺序运行的检测引擎，用于在用户输入到达 GenAI 应用程序之前对其进行分析和分类。
- 该系统支持 [Roota](https://github.com/UncoderIO/Roota) 和 [Sigma rules](https://sigmahq.io/docs/guide/about.html)，能够应用来自多个来源的检测逻辑，例如 [SigmaHQ](https://github.com/SigmaHQ/sigma)（发布时约有 1,200 条兼容的免费社区 Sigma 规则）、[SOC Prime](https://tdm.socprime.com/)（多达 3,000 条额外的兼容规则）以及其他第三方仓库。Sigma 规则可用于检测恶意软件利用本地 LLM 生成恶意代码以供执行的用例。
- SOC Prime [Uncoder AI](https://tdm.socprime.com/uncoder-ai/) 集成通过将 Sigma 规则转换为 Semgrep 格式，进一步扩展了功能，提供了标准化且可复用的检测 pipeline（需要免费账户）。
- Roota 规则为基于 regex 的 pipeline 提供支持。
- 该架构支持规则扩展性，可无缝集成组织特定的签名和外部检测内容。
- 该系统还可以作为本地日志传感器，记录用户和 agent 的 prompt，并支持诊断、事件发现和网络攻击调查。
- 检测逻辑与 [MITRE ATLAS](https://atlas.mitre.org/) 和 [OWASP Top 10 for LLMs](https://owasp.org/www-project-top-10-for-large-language-model-applications/) 等行业框架保持一致，确保针对对抗技术实现标准化覆盖。
- 操作包括允许、阻止或通知，具体取决于规则匹配和策略配置。
这种分层检测方法针对不断演变的对抗性 prompt 工程和其他 AI 聚焦的攻击向量提供了纵深防御。
灵感来源于 LlamaFirewall。
## 🚀 功能
- **多 Pipeline 检测**：Regex 模式、ML 模型、基于向量的相似度检测以及基于 LLM 的分析
- **灵活配置**：通过 JSON 进行动态 Pipeline 配置
- **实时分析**：具有可配置阈值的快速异步处理
- **客户端管理器**：灵活的客户端管理（Elasticsearch、OpenSearch、Qdrant）
- **RESTful API**：轻松与现有应用程序集成
- **可扩展架构**：用于自定义 Pipeline 的简单插件系统
## 🏗️ 架构
```
┌────────────────────────────────┐
│ 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) │ │
│ └──────────────────────────┘ │
└──────────────────────────────┘
```
## 📋 目录
- [安装说明](#%EF%B8%8F-installation)
- [配置](#%EF%B8%8F-configuration)
- [使用方法](#-usage)
- [API 参考](#-api-reference)
- [Pipelines](#-pipelines)
- [管理器](#-managers)
- [规则管理与定制](#-rule-management-and-customization)
- [添加自定义 Pipelines](#-adding-custom-pipelines)
- [开发](#-development)
- [许可证](#-license)
- [构建工具](#%EF%B8%8F-built-with)
- [TO-DO 列表](#%EF%B8%8F-to-do-list)
## 🛠️ 安装说明
### 前置条件
- Python 3.12+
- OpenSearch、Elasticsearch 或 Qdrant（通过 Similarity Manager）
- LLM API key（可选，用于 LLM Pipeline）：
- 用于 GPT 模型的 OpenAI API key，或
- 用于 Claude 模型的 Anthropic API key，或
- 用于企业部署的 Azure OpenAI 凭证，或
- 本地运行的 Ollama，用于注重隐私的本地 LLM 模型
### 快速开始
1. **克隆仓库**
git clone https://github.com/0xAIDR/AIDR-Bastion.git
cd AIDR-Bastion
2. **创建虚拟环境**
python -m venv venv
source venv/bin/activate # On Windows: venv\Scripts\activate
3. **安装依赖**
pip install -r requirements.txt
4. **设置环境变量**
cp env.example .env
# 使用您的配置编辑 .env
5. **运行应用程序**
python server.py
默认情况下，API 将在 `http://localhost:8000` 上可用。你可以使用 HOST 和 PORT 环境变量自定义主机和端口。
## ⚙️ 配置
### 环境变量 (.env)
```
# FastAPI 配置
HOST=0.0.0.0
PORT=8000
# ML Pipeline.
# 模型路径
ML_MODEL_PATH=
# LLM Pipeline
LLM_DEFAULT_CLIENT=openai # openai, anthropic, azure, or ollama
# OpenAI 配置
OPENAI_API_KEY=
OPENAI_MODEL=gpt-4
OPENAI_BASE_URL=https://api.openai.com/v1
# Anthropic 配置
ANTHROPIC_API_KEY=
ANTHROPIC_MODEL=claude-sonnet-4-5-20250929
ANTHROPIC_BASE_URL=https://api.anthropic.com
# Azure OpenAI 配置
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 配置 (用于本地 LLM 模型)
OLLAMA_BASE_URL=http://localhost:11434
OLLAMA_MODEL=llama3
# LLM 通用配置 (适用于所有 LLM 提供商)
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
SIMILARITY_PROMPT_INDEX=
SIMILARITY_NOTIFY_THRESHOLD=0.7
SIMILARITY_BLOCK_THRESHOLD=0.87
# Manager 配置
SIMILARITY_DEFAULT_CLIENT=opensearch # opensearch, elasticsearch, or qdrant
LLM_DEFAULT_CLIENT=openai # openai, anthropic, azure, or ollama
# OpenSearch 配置
OS__HOST=
OS__PORT=
OS__SCHEME=
OS__USER=
OS__PASSWORD=
# Elasticsearch 配置 (OpenSearch 的替代方案)
ES__HOST=
ES__PORT=
ES__SCHEME=
ES__USER=
ES__PASSWORD=
# Qdrant 配置 (OpenSearch/Elasticsearch 的替代方案)
QDRANT__HOST=localhost
QDRANT__PORT=6333
QDRANT__GRPC_PORT=6334
QDRANT__API_KEY=
QDRANT__PREFER_GRPC=false
QDRANT__TIMEOUT=30
# Kafka 配置 (用于事件日志记录)
KAFKA__BOOTSTRAP_SERVERS=localhost:9092
KAFKA__TOPIC=aidr-events
KAFKA__SECURITY_PROTOCOL=PLAINTEXT
KAFKA__SASL_MECHANISM=
KAFKA__SASL_USERNAME=
KAFKA__SASL_PASSWORD=
# 在 Similarity Pipeline 和 ML Pipeline 中创建嵌入所需
EMBEDDINGS_MODEL=
## Kafka 配置
# KAFKA__BOOTSTRAP_SERVERS=
# KAFKA__TOPIC=
# KAFKA__SECURITY_PROTOCOL=PLAINTEXT
# KAFKA__SASL_MECHANISM=
# KAFKA__SASL_USERNAME=
# KAFKA__SASL_PASSWORD=
# KAFKA__SAVE_PROMPT=true
```
### Pipeline 配置
`config.json` 文件控制每个流程将运行哪些 Pipelines。
默认 `config.json` 配置：
```
[
{
"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"
]
}
]
```
#### 配置影响
- **Flow names**：可以是任何自定义名称（例如 `base`、`code`、`security`、`content`）。该名称必须与你在 API 请求的 `pipeline_flow` 参数中传递的内容相匹配
- **Pipeline names**：必须与 `PipelineNames` 枚举中定义的 Pipeline 名称相匹配
- **顺序很重要**：Pipelines 按数组中指定的顺序运行
- **Example flows**：
- `base` flow：对一般文本 prompt 进行有害内容检测
- `code` flow：对代码片段进行安全漏洞检测
- `custom_flow`：你可以为特定用例创建任何自定义 flow 名称
## 🚀 使用方法
### 基本 API 用法
```
import requests
# 对文本提示运行 Pipeline 分析
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']}")
# 获取可用的 Flow 和 Pipeline
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']]}")
# 获取可用的 Manager 及其客户端
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']]}")
# 获取特定 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']}")
# 切换 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 用法
```
from app.main import bastion_app
# 直接使用
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}")
```
### 与现有应用程序集成
你可以将项目集成到现有的 LLM 应用程序中：
1. **发送请求：**
```
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. **配置你的应用程序以检查所有用户输入**
3. **设置适当的错误处理和回退机制**
4. **动态管理 managers 和 clients**：
```
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']
```
### 项目配置
项目可以通过环境变量进行配置：
- `HOST`：服务器主机（默认值：0.0.0.0）
- `PORT`：服务器端口（默认值：8000）
- `CORS_ORIGINS`：CORS 允许的源
- `EMBEDDINGS_MODEL`：用于 embeddings 的 Hugging Face 模型
- `SIMILARITY_NOTIFY_THRESHOLD`：通知阈值
- `SIMILARITY_BLOCK_THRESHOLD`：阻止阈值
你可以在 env.example 中找到所有必需的环境变量
## 📚 API 参考
### POST /api/v1/flow/run
运行 pipelines 以分析输入 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
获取所有可用 flow 及其 pipelines 的列表。
**Response:**
```
{
"flows": [
{
"flow_name": "string",
"pipelines": [
{
"name": "string",
"enabled": "boolean"
}
]
}
]
}
```
### GET /api/v1/manager/list
获取所有可用 manager 及其 client 的列表。
**Response:**
```
{
"managers": [
{
"id": "string",
"name": "string",
"enabled": "boolean",
"clients": ["string"]
}
]
}
```
### GET /api/v1/manager/{manager_id}
获取有关特定 manager 的信息。
**Parameters:**
- `manager_id` (string): manager 的 ID（例如 "similarity", "llm"）
**Response:**
```
{
"id": "string",
"name": "string",
"enabled": "boolean",
"clients": ["string"]
}
```
### POST /api/v1/manager/switch_active_client
切换特定 manager 的活动 client。
**Request Body:**
```
{
"manager_id": "string",
"client_id": "string"
}
```
**Response:**
```
{
"client_id": "string",
"status": "boolean"
}
```
## 🔍 Pipelines
### 1. Rule Pipeline (`rule`)
- **用途**：使用正则表达式进行基于模式的检测
- **规则**：位于 `app/pipelines/rule_pipeline/rules/` 的 YAML 文件
- **Categories**:
- **Injection**：SQL、命令执行、路径遍历
- **Obfuscation**：字符混淆、编码、Unicode 同形字
- **Override**：角色扮演、禁用过滤器、上下文拼接
- **Leakage**：直接 prompt 请求、强制重复
- **PII**：电子邮件、电话、信用卡、密码、API keys、UUIDs、IBAN
- **Semantic**：情感操纵、权威谬误、多语言攻击
- **DoS**：字符/单词重复、regex DoS
- **最适用于**：已知攻击模式和简单文本分析
### 2. Similarity Pipeline (`similarity`)
- **用途**：针对已知有害 prompt 进行基于向量的相似度检测
- **Backend**：使用 [Similarity Manager](#similarity-manager) 进行向量搜索
- **Required**：通过 [Similarity Manager](#similarity-manager) 配置 OpenSearch、Elasticsearch 或 Qdrant
- **Configuration**：`SIMILARITY_NOTIFY_THRESHOLD`、`SIMILARITY_BLOCK_THRESHOLD`、`SIMILARITY_DEFAULT_CLIENT`
- **最适用于**：检测已知攻击的变体
### 3. Code Analysis Pipeline (`code_analysis`)
- **用途**：使用 Semgrep 进行静态代码分析
- **Languages**：Python、JavaScript、Java、C++ 等
- **规则**：安全聚焦的模式
- **最适用于**：代码注入和漏洞检测
### 4. ML Pipeline (`ml`)
- **用途**：基于机器学习的分类
- **Configuration**：需要 `ML_PIPELINE_PATH`
- **模型**：用于 prompt 分类的自定义训练模型
- **最适用于**：一般恶意内容检测
- **Required**：已配置环境 `EMBEDDINGS_MODEL`
### 5. LLM Pipeline (`llm`)
- **用途**：使用高级语言模型进行 AI 驱动的分析
- **Backend**：使用 [LLM Manager](#llm-manager) 进行文本分析
- **支持的提供商**：
- **OpenAI** - GPT-4, GPT-4 Turbo 模型
- **Anthropic** - Claude 3.5 Sonnet, Claude 3 Opus 模型
- **Azure OpenAI** - 通过 Microsoft Azure 基础设施的企业级 GPT 模型
- **Ollama** - 本地 LLM 模型（Llama 3, Mistral, Gemma 等），适用于注重隐私的部署
- **Configuration**:
- `LLM_DEFAULT_CLIENT` - 选择提供商
- `LLM_TEMPERATURE` - 控制响应随机性（0.0-2.0，默认值：0.1）
- `LLM_MAX_TOKENS` - 最大响应长度（默认值：1000）
- 特定于提供商的 API keys 和设置
- **Features**：JSON 响应格式、可配置模型、智能决策制定、多提供商支持
- **Response Format**：返回结构化 JSON，包含状态和推理
- **最适用于**：复杂推理、上下文感知分析和细致入微的内容审核
## 👥 Managers
系统使用 manager 来控制不同类型的 client，并为 pipeline 提供统一的接口。
**可用的 managers：**
- **Similarity Manager** - 管理用于向量搜索的搜索 client
- **LLM Manager** - 管理用于文本分析的 LLM client
### Similarity Manager
管理用于相似内容向量搜索的搜索系统。根据配置自动选择可用的 client。
**可用的 clients：**
- **OpenSearch Client** - 带有 KNN 插件的全文本搜索，用于向量相似度
- **Elasticsearch Client** - 带有 dense_vector 支持的搜索引擎，用于余弦相似度
- **Qdrant Client** - 专为高性能相似度搜索优化的专用向量数据库
使用 SIMILARITY_DEFAULT_CLIENT 环境变量设置默认 client。
端点 `/api/v1/manager/switch_active_client` 也支持此操作。
**Qdrant 的优势：**
- 使用原生 HNSW 算法，向量搜索速度提升 2-3 倍
- 更低的内存占用和更好的资源效率
- 更简洁的 API，无需复杂的 KNN 查询
- 内置分数阈值过滤
- 原生支持 gRPC 协议
- 非常适合高吞吐量的生产环境
**开发中的 Clients：**
- 计划支持其他向量数据库
- 你可以贡献 client
### LLM Manager
管理用于文本分析和分类的 LLM 提供商。为多个 AI 提供商提供统一接口，支持动态 client 切换。
**可用的 clients：**
- **OpenAI Client** - GPT-4, GPT-4 Turbo 模型支持
- **Anthropic Client** - Claude 3.5 Sonnet, Claude 3 Opus 模型支持
- **Azure OpenAI Client** - 通过 Microsoft Azure 基础设施的企业级 GPT 模型
- **Ollama Client** - 本地 LLM 模型（Llama 3, Mistral, Gemma 等），适用于注重隐私的部署
使用 `LLM_DEFAULT_CLIENT` 环境变量设置默认 client。
端点 `/api/v1/manager/switch_active_client` 允许在提供商之间动态切换。
**切换提供商示例：**
```
# 切换到 Anthropic Claude
requests.post("http://localhost:8000/api/v1/manager/switch_active_client", json={
"manager_id": "llm",
"client_id": "anthropic"
})
# 切换到 OpenAI GPT
requests.post("http://localhost:8000/api/v1/manager/switch_active_client", json={
"manager_id": "llm",
"client_id": "openai"
})
# 切换到 Azure OpenAI
requests.post("http://localhost:8000/api/v1/manager/switch_active_client", json={
"manager_id": "llm",
"client_id": "azure"
})
# 切换到 Ollama (本地)
requests.post("http://localhost:8000/api/v1/manager/switch_active_client", json={
"manager_id": "llm",
"client_id": "ollama"
})
```
**开发中的 Clients：**
- 计划支持其他 LLM 提供商
- 你可以贡献 client
## 📋 规则管理与定制
### Rule Pipeline 的 YAML 规则
Rule Pipeline 使用 [Roota](https://github.com/UncoderIO/Roota) 规则文件定义检测模式。Roota 是一种用于集体网络防御的公共领域语言，它将 SIEM、EDR、XDR 或 Data Lake 的原生查询与标准化的元数据和威胁情报相结合，以实现自动翻译成其他语言。
每个规则文件遵循特定的结构
```
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 注入、命令执行、路径遍历、脚本注入
- **Obfuscation**：字符混淆、编码技巧、Unicode 同形字
- **Override**：角色扮演攻击、禁用过滤器、上下文拼接
- **Leakage**：直接 prompt 请求、强制重复攻击
- **PII**：个人身份信息检测（电子邮件、电话、信用卡等）
- **Semantic**：情感操纵、权威谬误、多语言攻击
- **DoS**：拒绝服务模式（字符重复、regex DoS）
### Code Analysis Pipeline 的 Semgrep 规则
Code Pipeline 使用 Semgrep 规则进行静态代码分析。规则位于 `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
```
### 规则管理
#### 使用 Roota 创建规则
[Roota](https://github.com/UncoderIO/Roota) 是一种用于集体网络防御的公共领域语言，它提供：
- **基于 YAML 的格式**，易于编写且人类可读
- **多语言支持**，适用于 SIEM、EDR、XDR 和 Data Lake 查询
- **MITRE ATT&CK 映射**，用于威胁情报集成
- **威胁行为者时间线**支持，用于协调防御
- **关联支持**，用于更强大的检测逻辑
- **OCSF 和 Sigma 兼容性**，以实现最大兼容性
**Roota 规则示例：**
```
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
```
#### 使用 Uncoder AI 生成 Semgrep 规则
[Uncoder AI](https://tdm.socprime.com/uncoder-ai/) 是一个强大的工具，用于将 Sigma 和 Roota 规则转换为各种格式，包括 Semgrep：
1. **访问 [Uncoder AI](https://tdm.socprime.com/uncoder-ai/)**
2. **免费注册账户**
3. **选择 Roota/Sigma 到 Semgrep 的转换**
4. **粘贴你的 Roota 或 Sigma 规则**
5. **生成 Semgrep YAML 规则**
6. **保存生成的规则**在 `app/pipelines/semgrep_pipeline/rules/` 中
**Sigma 规则示例：**
```
title: Suspicious PowerShell Command
description: Detects suspicious PowerShell commands
logsource:
category: process_creation
product: windows
detection:
selection:
- CommandLine: '*powershell*'
- CommandLine: '*Invoke-Expression*'
condition: selection
```
#### 使用 SOC Prime 获取高级规则
[SOC Prime](https://tdm.socprime.com/) 提供全面的威胁检测规则，包括 Roota 和 Sigma 格式：
1. **访问 [SOC Prime](https://tdm.socprime.com/)**
2. **浏览 Rules Library**（Sigma、Roota 和其他格式）
3. **按技术和威胁类型过滤**
4. **下载或使用 [Uncoder AI](https://tdm.socprime.com/uncoder-ai/) 转换规则**
5. **根据你的特定用例调整规则**
### 自定义规则开发
#### 创建自定义 Regex 规则
1. **识别攻击模式**
2. **在适当的类别文件夹中创建 YAML 文件**
3. **定义带有清晰描述的模式**
4. **使用各种输入进行彻底测试**
5. **设置适当的严重性和操作**
**自定义规则示例：**
```
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
```
#### 创建自定义 Semgrep 规则
1. **识别代码漏洞模式**
2. **使用 Semgrep 语法编写模式**
3. **使用示例代码进行测试**
4. **添加适当的元数据**
5. **放置在规则目录中**
**自定义 Semgrep 规则示例：**
```
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 选择
### 优先 Client 选择
- **Similarity Manager**：使用 `SIMILARITY_DEFAULT_CLIENT` 在 OpenSearch、Elasticsearch 或 Qdrant 之间进行选择
- **LLM Manager**：使用 `LLM_DEFAULT_CLIENT` 选择 LLM 提供商
- **动态切换**：可以通过 API 端点 `/api/v1/manager/switch_active_client` 更改活动 client
#### 优先级配置
```
# Similarity Manager 优先级
SIMILARITY_DEFAULT_CLIENT=opensearch # opensearch, elasticsearch, or qdrant
# LLM Manager 优先级
LLM_DEFAULT_CLIENT=openai # openai, anthropic, azure, or ollama
```
## ⚙️ 启用已禁用的 Pipeline
某些 pipeline 默认处于禁用状态。要启用它们：
### 使用 OpenAI 启用 LLM Pipeline
1. 在 .env 配置文件中配置你的 OpenAI API key：
LLM_DEFAULT_CLIENT=openai
OPENAI_API_KEY=your-api-key
OPENAI_MODEL=gpt-4
OPENAI_BASE_URL=https://api.openai.com/v1
# 可选：LLM 通用配置 (适用于所有提供商)
LLM_TEMPERATURE=0.1 # Lower = more focused, higher = more creative
LLM_MAX_TOKENS=1000 # Maximum response length
2. 在 `config.json` 中添加到你的 flow：
{
"flow_name": "base",
"pipelines": ["similarity", "rule", "llm"]
}
### 使用 Anthropic Claude 启用 LLM Pipeline
1. 在 .env 配置文件中配置你的 Anthropic API key：
LLM_DEFAULT_CLIENT=anthropic
ANTHROPIC_API_KEY=your-api-key
ANTHROPIC_MODEL=claude-sonnet-4-5-20250929
ANTHROPIC_BASE_URL=https://api.anthropic.com
# 可选：LLM 通用配置 (适用于所有提供商)
LLM_TEMPERATURE=0.1 # Lower = more focused, higher = more creative
LLM_MAX_TOKENS=1000 # Maximum response length
2. 在 `config.json` 中添加到你的 flow：
{
"flow_name": "base",
"pipelines": ["similarity", "rule", "llm"]
}
### 使用 Azure OpenAI 启用 LLM Pipeline
1. 在 .env 配置文件中配置你的 Azure OpenAI 凭证：
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
# 可选：LLM 通用配置 (适用于所有提供商)
LLM_TEMPERATURE=0.1 # Lower = more focused, higher = more creative
LLM_MAX_TOKENS=1000 # Maximum response length
2. 在 `config.json` 中添加到你的 flow：
{
"flow_name": "base",
"pipelines": ["similarity", "rule", "llm"]
}
### 使用 Ollama 启用 LLM Pipeline（本地模型）
1. 安装并启动 Ollama：
# 从 https://ollama.ai 安装 Ollama
# 拉取模型 (例如 Llama 3)
ollama pull llama3
# Ollama 自动在 http://localhost:11434 上启动
2. 在 .env 配置文件中配置 Ollama：
LLM_DEFAULT_CLIENT=ollama
OLLAMA_BASE_URL=http://localhost:11434
OLLAMA_MODEL=llama3
# 可选：LLM 通用配置 (适用于所有提供商)
LLM_TEMPERATURE=0.1 # Lower = more focused, higher = more creative
LLM_MAX_TOKENS=1000 # Maximum response length
3. 在 `config.json` 中添加到你的 flow：
{
"flow_name": "base",
"pipelines": ["similarity", "rule", "llm"]
}
4. 通过 API 动态切换提供商：
# 切换到 Anthropic
requests.post("http://localhost:8000/api/v1/manager/switch_active_client",
json={"manager_id": "llm", "client_id": "anthropic"})
# 切换到 OpenAI
requests.post("http://localhost:8000/api/v1/manager/switch_active_client",
json={"manager_id": "llm", "client_id": "openai"})
# 切换到 Azure OpenAI
requests.post("http://localhost:8000/api/v1/manager/switch_active_client",
json={"manager_id": "llm", "client_id": "azure"})
# 切换到 Ollama
requests.post("http://localhost:8000/api/v1/manager/switch_active_client",
json={"manager_id": "llm", "client_id": "ollama"})
### 启用 ML Pipeline
1. 在 .env 配置文件中配置模型路径：
ML_MODEL_PATH=/path/to/your/model
2. 在 `config.json` 中添加到你的 flow：
{
"flow_name": "base",
"pipelines": ["similarity", "rule", "ml"]
}
## 🔧 添加自定义 Pipelines
### 第 1 步：创建 Pipeline 类
```
# 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
)
```
### 第 2 步：注册 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
}
```
### 第 3 步：添加到配置
```
[
{
"flow_name": "base",
"pipelines": [
"personal_info",
"similarity",
"rule",
"my_pipeline"
]
}
]
```
### 第 4 步：将 Pipeline 名称添加到枚举
```
# app/core/enums.py
class PipelineNames(str, Enum):
# ... existing names
my_pipeline = "my_pipeline"
```
## 🧪 开发
### 设置 OpenSearch
1. **安装 OpenSearch**
docker run -p 9200:9200 -p 9600:9600 -e "discovery.type=single-node" opensearchproject/opensearch:latest
2. **创建 similarity index**
python app/scripts/similarity/index_script.py
这将在 OpenSearch 中创建 `similarity-prompt-index` 索引。你可以通过设置 SIMILARITY_PROMPT_INDEX 环境变量来自定义索引名称。
### 设置 Elasticsearch
1. **安装 Elasticsearch**
# OpenSearch 的替代方案
docker run -p 9200:9200 -e "discovery.type=single-node" elasticsearch:latest
2. **创建 similarity index**
python app/scripts/similarity/index_script.py
这将在 Elasticsearch 中创建 `similarity-prompt-index` 索引。你可以通过设置 SIMILARITY_PROMPT_INDEX 环境变量来自定义索引名称。
### 设置 Qdrant
1. **安装 Qdrant**
# 推荐：最快且最高效的选项
docker run -p 6333:6333 -p 6334:6334 qdrant/qdrant
2. **配置环境变量**
# 添加到您的 .env 文件
SIMILARITY_DEFAULT_CLIENT=qdrant
QDRANT__HOST=localhost
QDRANT__PORT=6333
QDRANT__GRPC_PORT=6334
# 可选：用于 Qdrant Cloud
# QDRANT__API_KEY=your-api-key
# QDRANT__PREFER_GRPC=true
3. **创建 similarity collection**
python scripts/similarity/index_script.py
这将在 Qdrant 中创建 `similarity-prompt-index` collection，并采用优化的向量搜索配置（HNSW 算法，余弦相似度）。
**为什么选择 Qdrant：**
- **性能**：向量操作速度比 OpenSearch/Elasticsearch 快 2-3 倍
- **效率**：通过优化的向量存储降低内存使用
- **简洁性**：更简洁的 API，无需复杂的查询 DSL
- **功能**：内置 payload 过滤、量化和快照
- **扩展性**：生产就绪，支持水平扩展和云选项
### 设置 Kafka 用于事件日志记录
AIDR Bastion 支持 Kafka 集成，用于记录 BLOCK 和 NOTIFY 事件，实现可扩展的事件流和实时监控。
#### 使用 Docker Compose 快速开始
**配置环境变量**
最少必需的环境变量
```
# 添加到您的 .env 文件
KAFKA__BOOTSTRAP_SERVERS=localhost:9092
KAFKA__TOPIC=aidr-events
KAFKA__SECURITY_PROTOCOL=PLAINTEXT
```
完整的 Kafka 环境变量
```
## Kafka 配置
# KAFKA__BOOTSTRAP_SERVERS=
# KAFKA__TOPIC=
# KAFKA__SECURITY_PROTOCOL=PLAINTEXT
# KAFKA__SASL_MECHANISM=
# KAFKA__SASL_USERNAME=
# KAFKA__SASL_PASSWORD=
# KAFKA__SAVE_PROMPT=true
```
环境变量 `KAFKA__SAVE_PROMPT` 是可选的。它控制是否将输入 prompt 数据保存到 Kafka。
#### 事件日志记录功能
- **BLOCK Events**：当 prompt 被检测规则阻止时记录
- **NOTIFY Events**：当 prompt 触发通知但被允许时记录
- **Structured JSON**：事件包含 prompt 内容、检测结果和元数据
- **Real-time Streaming**：事件立即发送到 Kafka 主题
#### 事件 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
}
```
#### 高级 Kafka 配置
对于生产环境，配置额外的安全设置：
```
# SASL 认证
KAFKA__SECURITY_PROTOCOL=SASL_SSL
KAFKA__SASL_MECHANISM=PLAIN
KAFKA__SASL_USERNAME=your_username
KAFKA__SASL_PASSWORD=your_password
# SSL 配置 (如果需要)
KAFKA__SSL_CA_LOCATION=/path/to/ca-cert
KAFKA__SSL_CERTIFICATE_LOCATION=/path/to/client-cert
KAFKA__SSL_KEY_LOCATION=/path/to/client-key
```
#### 规则管理
- **test_rules.py**：全面的规则测试和验证
- **generate_rules.py**：交互式规则创建和转换
- **[Roota](https://github.com/UncoderIO/Roota)**：用于集体网络防御的公共领域语言
- **[Uncoder AI](https://tdm.socprime.com/uncoder-ai/)**：将 Roota/Sigma 规则转换为 Semgrep 格式
- **[SOC Prime](https://socprime.com/)**：访问全面的威胁检测规则
## 📄 许可证
本项目采用 GNU Lesser General Public License v3.0 (LGPL-3.0) 许可 - 详情请参阅 [LICENSE](LICENSE) 文件。
有关 LGPL 的更多信息，请访问：https://www.gnu.org/licenses/lgpl-3.0.html
## 🛠️ 构建工具
本项目使用以下强大的开源库和框架构建：
- **[Roota](https://github.com/UncoderIO/Roota)** - 用于集体网络防御的公共领域语言
- **[Uncoder AI](https://tdm.socprime.com/uncoder-ai/)** - 将 Roota/Sigma 规则转换为 Semgrep 格式
- **[FastAPI](https://fastapi.tiangolo.com/)** - 现代、快速的 Web 框架，用于基于标准 Python 类型提示构建 API（Python 3.7+）
- **[OpenSearch](https://opensearch.org/)** - 开源搜索和分析套件，用于日志分析、应用搜索等
- **[Elasticsearch](https://www.elastic.co/elasticsearch/)** - 开放搜索和分析平台，适用于各种数据类型
- **[OpenAI](
标签：AMSI绕过, AV绕过, DNS 反向解析, FastAPI, GenAI安全, LLM防火墙, MITRE ATLAS, OWASP Top 10 LLM, Python, Roota, Semgrep, Sigma规则, SOC工具, WordPress安全扫描, 云计算, 人工智能安全, 内容安全审核, 合规性, 大模型防护, 威胁检测, 恶意代码生成检测, 恶意输入检测, 提示词注入防护, 无后门, 漏洞发现, 目标导入, 网络安全, 规则引擎, 隐私保护