icepaule/IcePorge-Malware-RAG
GitHub: icepaule/IcePorge-Malware-RAG
基于 RAG 与本地 LLM 的恶意软件分析问答系统，将专业知识库与语义检索结合，为安全分析提供上下文丰富的 AI 辅助。
Stars: 1 | Forks: 0
# IcePorge 恶意软件 RAG
**结合检索增强生成的 AI 驱动恶意软件分析**
[IcePorge](https://github.com/icepaule/IcePorge) 恶意软件分析栈的一部分。
[](https://opensource.org/licenses/MIT)
## 概述
Malware RAG 通过将精选的知识库（SANS FOR610/SEC504）与基于 Ollama 的本地 LLM 推理相结合，为恶意软件分析提供上下文丰富的 AI 响应。它使用 ChromaDB 进行向量存储和语义搜索。
### 核心功能
- **RAG Pipeline** - 用于准确响应的检索增强生成
- **FOR610 知识库** - SANS 恶意软件分析课程资料
- **本地 LLM** - 通过 Ollama 进行注重隐私的推理（无云 API）
- **向量搜索** - 基于 ChromaDB 的语义相似度搜索
- **REST API** - 易于与分析工具集成
## 架构
```
+--------------------------------------------------+
| Malware RAG (ki01) |
| |
| +--------------------------------------------+ |
| | rag_api.py | |
| | Flask REST API (:5001) | |
| +--------------------------------------------+ |
| | | |
| v v |
| +-------------+ +-------------+ |
| | ChromaDB | | Ollama | |
| | (Vectors) | | (qwen2.5) | |
| +-------------+ +-------------+ |
| ^ |
| | |
| +-------------+ |
| | FOR610/SEC | |
| | 504 PDFs | |
| +-------------+ |
+--------------------------------------------------+
Query Flow:
1. User query received
2. Query embedded via SentenceTransformer
3. Similar chunks retrieved from ChromaDB
4. Context + Query sent to Ollama
5. Enriched response returned
```
## 组件
| 文件 | 描述 |
|------|-------------|
| `rag_api.py` | 包含 RAG pipeline 的 Flask REST API |
| `ingest_documents.py` | PDF 摄取与分块 |
## 安装说明
### 前置条件
- Python 3.10+
- 包含 `qwen2.5-coder:14b` 模型的 Ollama
- 16GB+ RAM（用于 embeddings 和 LLM）
### 设置
```
# 克隆 repository
git clone https://github.com/icepaule/IcePorge-Malware-RAG.git
cd IcePorge-Malware-RAG
# 创建虚拟环境
python3 -m venv venv
source venv/bin/activate
# 安装依赖
pip install flask flask-cors chromadb sentence-transformers requests langchain langchain-community pypdf
# 导入文档（将 PDF 放入 documents/ 文件夹）
python ingest_documents.py /path/to/FOR610.pdf /path/to/SEC504.pdf
# 启动 API 服务器
python rag_api.py
```
## API 端点
### 健康检查
```
curl http://localhost:5001/health
```
响应：
```
{
"status": "healthy",
"service": "malware-rag",
"collections": ["for610"],
"ollama_available": true
}
```
### 结合 RAG 上下文的查询
```
curl -X POST http://localhost:5001/query \
-H "Content-Type: application/json" \
-d '{"query": "How do I identify process injection techniques?", "collection": "for610"}'
```
响应：
```
{
"response": "Process injection can be identified by...",
"sources": [
{"content": "...", "source": "FOR610-book1.pdf", "page": 45, "relevance": 0.89}
],
"model": "qwen2.5-coder:14b"
}
```
### 直接 LLM 查询（无 RAG）
```
curl -X POST http://localhost:5001/llm \
-H "Content-Type: application/json" \
-d '{"prompt": "Explain RC4 encryption in malware"}'
```
## 文档摄取
### 支持的格式
- PDF（通过 PyPDF）
- 文本文件
### 摄取流程
```
# 单个文档
python ingest_documents.py /path/to/document.pdf
# 多个文档
python ingest_documents.py /path/to/doc1.pdf /path/to/doc2.pdf
# 自定义 collection 名称
python ingest_documents.py --collection sec504 /path/to/SEC504.pdf
```
### 分块配置
| 参数 | 默认值 | 描述 |
|-----------|---------|-------------|
| `chunk_size` | 1000 | 每个 chunk 的字符数 |
| `chunk_overlap` | 200 | chunk 之间的重叠量 |
| `embedding_model` | `all-MiniLM-L6-v2` | SentenceTransformer 模型 |
## 配置
### 环境变量
| 变量 | 默认值 | 描述 |
|----------|---------|-------------|
| `CHROMA_PATH` | `/opt/malware-rag/chroma_db` | ChromaDB 存储路径 |
| `OLLAMA_URL` | `http://localhost:11434` | Ollama API endpoint |
| `DEFAULT_MODEL` | `qwen2.5-coder:14b` | 默认 LLM 模型 |
## 集成示例
### Python 客户端
```
import requests
def analyze_with_context(query):
response = requests.post(
"http://ki01:5001/query",
json={"query": query, "collection": "for610"}
)
return response.json()
# 示例：获取上下文感知分析
result = analyze_with_context("What APIs are commonly used for keylogging?")
print(result["response"])
```
### 命令行
```
# 快速查询
curl -s http://ki01:5001/query \
-H "Content-Type: application/json" \
-d '{"query": "Explain API hashing"}' | jq .response
```
## 知识库集合
| 集合 | 内容 | 文档 |
|------------|---------|-----------|
| `for610` | SANS FOR610 - 恶意软件逆向工程 | 课程 PDF |
| `sec504` | SANS SEC504 - 事件响应 | 课程 PDF |
## 服务管理
```
# 作为 systemd 服务启动
sudo systemctl start malware-rag
# 查看日志
sudo journalctl -u malware-rag -f
# 检查状态
sudo systemctl status malware-rag
```
## 性能调优
### GPU 加速
为了获得更快的 LLM 推理速度，请确保 Ollama 使用 GPU：
```
# 检查 GPU 使用情况
nvidia-smi
# Ollama 会在可用时自动使用 GPU
ollama run qwen2.5-coder:14b
```
### 内存优化
- 减少查询中的 `n_results` 以获取更少的上下文
- 使用更小的 embedding 模型以加快搜索速度
- 调整 Ollama 中的 `num_ctx` 以改变上下文窗口大小
## 许可证
MIT 许可证 - 参见 [LICENSE](LICENSE)
**作者：** Michael Pauli
- GitHub: [@icepaule](https://github.com/icepaule)
- 邮箱：info@mpauli.de
标签：AI风险缓解, DAST, DLL 劫持, Ruby, 向量数据库, 大语言模型, 恶意软件分析, 本地推理, 检索增强生成, 知识库, 逆向工具