ties2/BBAP-Sec-AI-Malware-Detector
GitHub: ties2/BBAP-Sec-AI-Malware-Detector
一个结合静态 ML、规则分析与 LLM 的反编译驱动恶意软件检测流水线，旨在通过混合分层架构解决传统静态模型应对新型威胁时性能下降的问题。
Stars: 0 | Forks: 0
# BBAP-Sec AI 恶意软件检测器
**一个结合静态 ML、基于规则的代码分析和 LLM 的反编译驱动恶意软件检测 pipeline。**
这是论文*"A Decompilation-Driven Framework for Malware Detection with Large Language Models"* (Chawla & Prasad, 2026) 中方法的实现，并打包了生产级的 MLOps 架构（自定义日志、MLflow、FastAPI、Docker、CI/CD）。
它通过结合三个互补的检测器，将 Windows 可执行文件分类为**恶意软件** (malware) 或**良性** (benign)，并设计为可在**本地机器上轻量级运行**（无需 GPU，无需强制安装 Ghidra，无需强制 API key）。
## 为什么这样设计
论文的核心发现驱动了这里的每一个架构选择：
| 论文发现 | 本项目的应对方案 |
|---|---|
| 传统 ML (XGBoost) 对*已知*恶意软件的检测非常准确 (~98.5%) | 快速的**静态 XGBoost** 分支优先在本地运行。 |
| ...但对*新*威胁的检测性能急剧下降 (降至 ~74%) | **混合升级**：不确定的情况会移交给 LLM 分支。请参阅 `make eval-drift` 查看关于此性能下降的实时演示。 |
| LLM 能增加价值，但速度慢/成本高/需要 API | LLM 分支是**可选的**，仅在不确定的情况下调用。 |
| 反编译的 C 代码会溢出 context window（"迷失在中间"） | **显著性排序的分块** (Salience-ranked chunking) 将最安全相关的函数保留在窗口内。 |
| LLM 的决策是黑盒（可解释性差距） | **基于规则的类 YARA 扫描器**生成透明的、针对每个指标的推理，并且 XGBoost 会报告主要的特征驱动因素。 |
| 模型需要持续重新训练 | **PSI drift 监控**会在实际数据分布发生变化时发出警报，提示需要重新训练。 |
## 架构
```
┌─────────────────────────────────────────────┐
binary (.exe/.dll) ──► │ Stage 1: Quarantine (SHA256 + isolate) │
│ [optional] Ghidra headless decompile│──► decompiled C
└─────────────────────────────────────────────┘
│ │
┌────────────────────┘ └───────────────────┐
▼ ▼
┌──────────────────────────┐ ┌────────────────────────────────────┐
│ Stage 2a: STATIC arm │ │ Stage 2b/2c: CODE arms │
│ pefile → feature vector │ │ • Heuristic scanner (offline) │
│ → XGBoost P(malware) │ │ • LLM classifier (optional, API) │
└──────────────────────────┘ └────────────────────────────────────┘
│ │
└────────────────────────────┬───────────────────────────────────-─┘
▼
┌─────────────────────────────────────────────┐
│ Stage 3: Fusion → verdict + confidence │
│ + rationale (explainability) │
└─────────────────────────────────────────────┘
```
**三个检测分支：**
1. **静态 (XGBoost)** — 使用 `pefile` 提取 PE header/section/import 特征（绝不执行文件）并进行分类。快速，完全离线。
2. **启发式扫描器** — 基于规则，对反编译的 C 代码进行 YARA 风格的检测。离线，完全可解释，无需 API key。
3. **LLM 分类器** — 使用论文中精确的 zero-shot prompt，将反编译的 C 代码发送给 Claude 或 Gemini。可选；针对困难情况的智能升级。
**Pipeline 模式** (`config.yaml → pipeline.mode`)：
- `static_only` — 仅使用 XGBoost。最快。
- `llm_only` — 仅使用代码分支。
- `hybrid` *(默认)* — 优先使用 XGBoost；仅将不确定的情况（`P(malware)` 处于升级区间内）升级给 LLM。
## 快速开始（本地，约 2 分钟）
```
# 1. 安装（仅需轻量级核心 deps）
python -m venv .venv && source .venv/bin/activate # Windows: .venv\Scripts\activate
make setup # or: pip install -r requirements.txt && pip install -e .
# 2. 引导一个可工作的 model + 运行 end-to-end demo（无需 API key，无需 Ghidra）
make demo
```
`make demo` 会在合成数据上训练静态模型，然后对内置的样本进行分类：
- `data/samples/benign_sample.c` → **benign**
- `data/samples/suspicious_sample.c` → **malware** *(匹配项: process_injection, anti_debug_evasion, dynamic_api_resolution, persistence)*
### 运行 API
```
make serve # http://localhost:8000/docs (interactive Swagger UI)
```
```
# 分类粘贴的已反编译 C 代码（无需 binary）：
curl -X POST http://localhost:8000/analyze/code \
-H "Content-Type: application/json" \
-d '{"code": "void f(){ OpenProcess(); VirtualAllocEx(); CreateRemoteThread(); }", "file_name": "test.c"}'
# 上传真实的 binary（运行包含 static features 的完整 pipeline）：
curl -X POST http://localhost:8000/analyze/file -F "file=@/path/to/sample.exe"
# Health + capability 报告：
curl http://localhost:8000/health
```
## 常用命令
| 命令 | 作用 |
|---|---|
| `make demo` | 初始化模型 + 运行端到端 demo |
| `make train` | 训练静态模型（受 MLflow 跟踪） |
| `make evaluate` | 在全新的拆分数据集上进行评估 |
| `make eval-drift` | **演示论文的发现**：在模拟的较新威胁上准确率/召回率下降 |
| `make serve` | 运行 FastAPI 服务器 |
| `make test` | 运行 pytest 测试套件（21 个测试） |
| `make lint` | Ruff 检查 + 格式化 |
| `make mlflow` | 在 :5000 端口启动 MLflow UI |
| `make docker` / `make docker-run` | 构建 / 运行服务容器 |
### 漂移演示（适合在面试中展示）
```
make evaluate # baseline ~99.7% accuracy on known-distribution malware
make eval-drift # ~85% accuracy, recall collapses to ~70% on "newer" malware
```
这重现了论文的核心论点，即静态模型在面对不断演变的威胁时显得脆弱——这正是混合 LLM 升级和 drift 监控存在的原因。
## 启用 LLM 分支（可选）
```
cp .env.example .env # add ANTHROPIC_API_KEY or GOOGLE_API_KEY
pip install anthropic # or: pip install google-generativeai
```
在 `config/config.yaml` 中：
```
models:
llm:
enabled: true
provider: anthropic # or: google
```
LLM 分支使用了**论文中精确的 zero-shot prompt**（Listing 1），定义在 `src/models/llm_classifier.py` 中，并扩展了一行要求提供理由的请求，以提高分析师的透明度。
## 启用 Ghidra 反编译（可选，较重）
Ghidra 不是一个 pip 包并且需要依赖 JVM，因此它**默认禁用**。要启用对二进制文件的真实反编译：
1. 从 下载 Ghidra 并解压（需要 JDK 17+）。
2. 在 `config/config.yaml` 中：
decompilation:
enabled: true
ghidra_home: /opt/ghidra_11.x_PUBLIC # 你的安装路径
封装脚本 (`src/decompilation/ghidra_runner.py`) 会调用 `analyzeHeadless` 并运行后处理脚本 `src/decompilation/ghidra_scripts/decompile.py`，该脚本会将每个函数反编译为 C 代码。如果 Ghidra 不可用，pipeline 会优雅地降级到静态 + 启发式分支。
## 在真实数据上训练（Colab / Kaggle）
要获得一个达到作品集级别的模型，请使用真实数据集而不是合成数据进行训练：
1. 在 Google Colab 或 Kaggle 中打开 `notebooks/train_model_colab.ipynb`。
2. 将 `DATA_SOURCE` 设置为 `'ember'`（标准的 110 万样本 EMBER 基准测试）或 `'csv'`（任何包含 PE 特征的 Kaggle 数据集）。
3. 运行所有单元格；下载 `xgboost_malware.json` + `xgboost_malware.meta.json`。
4. 将这两个文件放入 `models/` — pipeline 和 API 会自动使用它们。
该 notebook 还记录了如何完全按照论文所述构建 JSONL 并运行 **Vertex AI LLM fine-tuning**（Gemini，40 个 epoch）。
## 项目结构
```
ai-malware-detector/
├── config/config.yaml # all tunables (paths, thresholds, model choices)
├── src/
│ ├── utils/ # logger (per-component audit logs) + config loader
│ ├── decompilation/
│ │ ├── ghidra_runner.py # headless Ghidra wrapper (Stage 1, optional)
│ │ └── ghidra_scripts/decompile.py # runs inside Ghidra's JVM
│ ├── features/
│ │ ├── schema.py # canonical feature order + suspicious-API list
│ │ ├── static_features.py # pefile-based PE feature extraction
│ │ ├── code_features.py # salience ranking + context-window chunking
│ │ └── code_heuristics.py # rule-based YARA-like scanner (explainable)
│ ├── models/
│ │ ├── xgboost_classifier.py # static model wrapper (save/load/explain)
│ │ ├── llm_classifier.py # LLM arm (paper's prompt, multi-provider)
│ │ ├── train.py # MLflow-tracked training + synthetic data
│ │ └── evaluate.py # metrics tables + drift simulation
│ ├── pipeline/detection_pipeline.py # orchestrator + verdict fusion
│ ├── serving/ # FastAPI app + Pydantic schemas
│ └── monitoring/drift_detector.py # PSI-based data/concept drift
├── scripts/
│ ├── generate_demo_model.py # one-shot bootstrap
│ └── run_demo.py # end-to-end demo
├── notebooks/train_model_colab.ipynb
├── tests/ # 21 pytest tests
├── deployment/ # Dockerfile + docker-compose (API + MLflow)
├── .github/workflows/ci-cd.yaml # lint → test → build image
├── Makefile pyproject.toml requirements.txt
```
## 样本的安全处理
本项目是一个**防御性**工具。它遵循论文中描述的安全处理实践：
- **静态优先**：静态和启发式分支绝不执行文件。
- **隔离**：上传的二进制文件会被计算哈希值 (SHA256) 并复制到隔离的 `data/quarantine/` 目录中，而不是在原位进行分析。
- 对于真实的恶意软件，仅在**没有网络连接的隔离 VM**内进行分析，之后恢复到干净的快照，并且永远不要将原始样本转移到通用系统中。
- 从信誉良好、获得许可的仓库（例如 MalwareBazaar，CC0）获取样本。本仓库**不包含真实的恶意软件** — `suspicious_sample.c` 是无实际功能的说明性伪代码。
内置样本和合成数据可以在任何地方安全运行。
## 局限性与未来工作
反映了论文中的讨论：
- **Context window** — 目前的分块是基于显著性的；层次化摘要将能够扩展到非常大的二进制文件。
- **可解释性** — 启发式分支和特征重要性有所帮助，但完整的 LLM 决策推理审计仍是未来的工作。
- **持续学习** — drift detection 会标记*何时*需要重新训练；自动化的重新训练触发器是自然的下一步。
- **端到端 LLM** — 论文指出 Ghidra 最终可能被 LLM 反编译器 (LLM4Decompile) 取代；`decompilation` 模块已被隔离，以便轻松进行这种替换。
- **多模态** — 添加动态/沙箱 trace 将推动实现生产级的可靠性。
## 技术栈
Python 3.9+ · XGBoost · scikit-learn · pefile · FastAPI · Pydantic · MLflow · Ghidra (可选) · Anthropic/Google LLM APIs (可选) · Docker · GitHub Actions · Ruff · pytest
标签：Apex, DLL 劫持, DNS 反向解析, MLOps, URL提取, 云安全监控, 反编译, 大语言模型, 安全规则引擎, 机器学习, 请求拦截, 逆向工具, 静态分析