Sushmit94/Securify
GitHub: Sushmit94/Securify
一个结合静态分析、微调机器学习分类器和 LLM 推理的源代码漏洞自动化检测框架，旨在比单一方法实现更高精度和更丰富的漏洞解释。
Stars: 0 | Forks: 0
# 基于AI的源代码漏洞自动化检测框架
这是一个混合多层系统，通过结合**静态分析**、**微调的机器学习分类器**和**基于LLM的推理**来检测源代码中的安全漏洞。
## 目录
1. [概述](#overview)
2. [动机](#motivation)
3. [系统设计](#system-design)
4. [架构层级详解](#architecture-layers-explained)
5. [语言支持策略](#language-support-strategy)
6. [技术栈](#tech-stack)
7. [文件夹结构](#folder-structure)
8. [数据集](#datasets)
9. [设置与安装](#setup--installation)
10. [用法](#usage)
11. [模型训练（微调）](#model-training-fine-tuning)
12. [评估指标](#evaluation-metrics)
13. [API参考](#api-reference)
14. [项目时间线](#project-timeline)
15. [局限性与未来工作](#limitations--future-work)
16. [参考](#references)
## 概述
本项目构建了一个框架，该框架将源代码（单个文件、多个文件或完整的 GitHub repository）作为输入，并生成结构化的漏洞报告，识别以下内容：
- **是什么**漏洞（例如：SQL注入、XSS、缓冲区溢出、硬编码凭据、不安全的反序列化）
- **在哪里**定位（文件、行号、函数）
- 系统对该发现的**置信度**
- **为什么**它是一个漏洞（通俗的英文解释）
- **如何修复**（建议的修复方案）
与纯基于规则的静态分析器或纯基于LLM的工具不同，该框架结合了三个独立的检测层并协调它们的输出，旨在实现**比任何单一方法更高的精度和更丰富的解释。**
## 动机
现有的方法通常分为两大阵营：
- **静态分析器**（Semgrep、Bandit、SonarQube）速度快且具有确定性，但它们是基于规则的——它们会遗漏新型或上下文相关的漏洞，并且经常产生误报。
- **纯基于LLM的工具**很灵活，能够推理上下文，但它们是不确定的，可能会产生幻觉，并且在对大型代码库进行分析时，如果没有事先过滤，运行成本会很高。
本项目结合了这两者，外加一个在标记的漏洞数据上训练的监督式机器学习分类器，以便：
- 静态分析提供快速、确定性的首轮检测。
- 机器学习分类器为存在标记数据集的语言添加了一个学习到的、数据驱动的信号。
- LLM推理层解释发现，建议修复方法，过滤误报，并将覆盖范围扩展到没有训练分类器的语言。
## 系统设计
```
┌─────────────────────────┐
│ Code Input (any lang) │
│ file / repo / snippet │
└───────────┬─────────────┘
│
┌───────────▼─────────────┐
│ Language Detector │
│ (extension + guesslang) │
└───────────┬─────────────┘
│
┌───────────▼─────────────┐
│ Language-specific │
│ Parser / Router │
└───────────┬─────────────┘
│
┌───────────────────────┼───────────────────────┐
│ │ │
┌───────▼────────┐ ┌─────────▼─────────┐ ┌─────────▼─────────┐
│ Layer 1: │ │ Layer 2: │ │ Layer 3: │
│ Static Analysis │ │ ML Classifier │ │ LLM Reasoning │
│ Semgrep ruleset │ │ (runs only if a │ │ (works regardless │
│ per language │ │ trained model │ │ of language) │
│ (30+ languages) │ │ exists for lang) │ │ │
└───────┬────────┘ └─────────┬─────────┘ └─────────┬─────────┘
│ (skipped for unsupported langs) │
└───────────────────────┼────────────────────────┘
│
┌───────────▼─────────────┐
│ Aggregator / Scorer │
│ merges findings, dedupes│
│ weights confidence based│
│ on which layers ran │
└───────────┬─────────────┘
│
┌───────────▼─────────────┐
│ Report Generator │
│ (JSON, HTML, PDF) │
└───────────┬─────────────┘
│
┌───────────▼─────────────┐
│ Web Dashboard (React) │
│ upload, view, export │
└─────────────────────────┘
```
### 数据流，逐步说明
1. **输入** — 用户上传文件、粘贴代码片段或提供 repo URL。
2. **语言检测器** — 通过扩展名识别每个文件的语言，对于模棱两可或无扩展名的文件，回退到轻量级分类器（`guesslang`）。
3. **解析器/路由器** — 将代码解析为 AST（如果存在该语言的解析器），并提取用于分析的函数/代码单元；将代码路由到相应的规则集和模型。
4. **第1层（静态分析）** — Semgrep（以及专门用于 Python 的 Bandit）针对特定语言的规则集进行扫描，并返回确定性的、基于模式的发现。
5. **第2层（ML分类器）** — 对于拥有训练模型的语言（最初是 Python、Java、C/C++），每个函数都会被标记化，并传递给微调过的 CodeBERT/GraphCodeBERT 模型，该模型输出一个漏洞概率得分。不支持的语言会跳过此步骤。
6. **第3层（LLM推理）** — 代码连同第1层和第2层的发现一起，通过结构化的 prompt 经由 API 发送给 LLM。LLM 确认/驳回发现，对其进行解释，并建议修复方法。此层适用于任何语言。
7. **聚合器** — 跨层合并并去重发现结果，分配最终的置信度得分（根据该语言实际运行的层级进行不同的加权），并按严重程度对结果进行排名。
8. **报告生成器** — 生成结构化的 JSON 报告以及人类可读的 HTML/PDF 版本。
9. **Dashboard** — React 前端，用于上传代码、浏览发现和导出报告。
## 架构层级详解
### 第1层 — 静态分析
- **工具：** Semgrep（主要工具，支持 30+ 种语言）、Bandit（专用于 Python，更深入的基于 AST 的检查）
- **输出：** 确定性的发现，包含规则 ID、文件、行号、严重程度、CWE 参考
- **角色：** 快速的第一轮扫描，捕获众所周知的漏洞模式
### 第2层 — ML分类器（微调）
- **基础模型：** CodeBERT 或 GraphCodeBERT（HuggingFace，约 125M 参数）
- **微调数据来源：** Devign、Big-Vul 和/或 SARD/Juliet Test Suite（带有标记的漏洞/安全函数对）
- **输出：** 每个函数的漏洞概率得分 (0–1)
- **角色：** 学习到的、数据驱动的信号，能够捕获静态规则遗漏的模式；仅适用于具有充足标记训练数据的语言
- **注意：** 这是唯一涉及实际模型训练（微调）的层；第1层和第3层按原样使用现有的工具/模型
### 第3层 — LLM推理
- **模型：** 通过 API 访问的现有大型语言模型（例如：Claude 或 GPT-4）— 未经过训练或微调
- **输入：** 代码片段 + 第1层的发现 + 第2层的得分（如果可用）
- **输出：** 包含已确认/驳回的发现、通俗易懂的英文解释和建议修复方案的结构化 JSON
- **角色：** 上下文推理、误报过滤，以及无论是否有 ML/静态工具支持都能覆盖任何语言
### 聚合器
- 对来自不同层中指向同一行/函数的重叠发现进行去重
- 根据**实际贡献的层级**对置信度进行加权 — 例如，由所有三层确认的发现，其权重要高于仅来自静态分析的发现
- 生成每个文件的最终漏洞排名列表
## 语言支持策略
由于第2层实际上无法针对每种语言进行训练（除了少数几种语言外，几乎没有现成的标记数据集），其覆盖范围被有意进行了分级：
| 级别 | 语言 | 启用的层级 |
|---|---|---|
| 第1级 (深度覆盖) | Python, Java, C/C++ | 静态分析 + ML分类器 + LLM推理 |
| 第2级 (广泛覆盖) | JavaScript/TypeScript, Go, PHP, Ruby, C#, Kotlin, Rust 和其他受 Semgrep 支持的语言 | 静态分析 + LLM推理 |
| 第3级 (回退) | 任何其他语言 | 仅 LLM推理（无静态规则集或训练模型） |
这种分级方法在报告输出中都有透明的记录——该框架从不会针对未训练过的语言声称具有基于机器学习的置信度。
## 技术栈
| 组件 | 技术 | 用途 |
|---|---|---|
| 静态分析 | Semgrep, Bandit | 基于规则的漏洞检测 |
| ML框架 | PyTorch, HuggingFace Transformers | 微调 CodeBERT/GraphCodeBERT |
| 数据集处理 | Pandas, HuggingFace Datasets | 加载/预处理 Devign, Big-Vul |
| LLM集成 | Claude API / OpenAI API | 推理、解释、修复建议 |
| 后端 | Python, FastAPI | 协调编排、REST API |
| 语言检测 | 文件扩展名映射 + `guesslang` | 将代码路由到正确的 pipeline |
| 代码解析 | `ast` (Python)、`esprima`/`@babel/parser` (JS)、`javalang` (Java)、`pycparser`/`clang` 绑定 (C/C++) | 提取用于分析的函数 |
| 前端 | React, Tailwind CSS | Dashboard UI |
| 数据库 | PostgreSQL | 扫描历史、发现结果存储 |
| 报告生成 | WeasyPrint / ReportLab | PDF 导出 |
| 实验追踪 | Weights & Biases (可选) | 记录微调运行 |
| 容器化 | Docker, Docker Compose | 可复现部署 |
| 测试 | Pytest | 单元和集成测试 |
## 文件夹结构
```
vuln-detection-framework/
├── README.md
├── requirements.txt
├── docker-compose.yml
├── .env.example
│
├── data/
│ ├── raw/
│ │ ├── python/
│ │ ├── java/
│ │ └── cpp/
│ ├── processed/
│ └── dataset_loader.py
│
├── language_support/
│ ├── detector.py # detects language of input file
│ ├── language_registry.py # maps language -> {ruleset, ml_model?, parser}
│ └── parsers/
│ ├── python_parser.py
│ ├── js_parser.py
│ ├── java_parser.py
│ ├── cpp_parser.py
│ └── generic_fallback.py
│
├── ml_model/
│ ├── train.py # fine-tuning script
│ ├── evaluate.py # precision/recall/F1 on test set
│ ├── model_config.py
│ ├── inference.py
│ └── checkpoints/
│ ├── python_model/
│ ├── java_model/
│ └── cpp_model/
│
├── static_analysis/
│ ├── semgrep_runner.py
│ ├── bandit_runner.py
│ └── rulesets/
│ ├── python.yml
│ ├── javascript.yml
│ ├── java.yml
│ ├── cpp.yml
│ └── go.yml
│
├── llm_reasoning/
│ ├── prompt_templates.py
│ ├── llm_client.py
│ └── response_parser.py
│
├── aggregator/
│ ├── merge_findings.py
│ └── scorer.py
│
├── backend/
│ ├── main.py
│ ├── routes/
│ │ ├── scan.py
│ │ └── report.py
│ ├── models/
│ └── services/
│ └── pipeline.py
│
├── frontend/
│ ├── src/
│ │ ├── components/
│ │ ├── pages/
│ │ └── App.jsx
│ └── package.json
│
├── reports/
│ └── generator.py
│
├── tests/
│ ├── test_static_analysis.py
│ ├── test_ml_model.py
│ ├── test_llm_reasoning.py
│ └── test_pipeline.py
│
└── docs/
├── system_design.md
├── language_coverage.md
├── dataset_notes.md
└── evaluation_results.md
```
## 数据集
| 数据集 | 语言 | 描述 |
|---|---|---|
| **Devign** | C/C++ | 来自真实开源项目的函数级别的漏洞/非漏洞标签 |
| **Big-Vul** | C/C++ | 带有 CVE 映射的大规模漏洞数据集 |
| **SARD / Juliet Test Suite** | C/C++, Java | 涵盖多种 CWE 类别的合成测试用例 |
| **CVEfixes** | 多种 | 将真实世界的 CVE 修复提交映射到漏洞/已修补代码对 |
这些用于为第1级语言微调第2层 ML 分类器。
## 设置与安装
```
# 克隆仓库
git clone
cd vuln-detection-framework
# 创建虚拟环境
python -m venv venv
source venv/bin/activate # or venv\Scripts\activate on Windows
# 安装依赖
pip install -r requirements.txt
# 安装 frontend 依赖
cd frontend && npm install && cd ..
# 设置环境变量
cp .env.example .env
# 将你的 LLM API key 和 database URL 添加到 .env
# 使用 Docker Compose 运行（推荐）
docker-compose up --build
```
## 用法
```
# 对单个文件运行扫描
python backend/services/pipeline.py --input path/to/file.py
# 对完整仓库运行扫描
python backend/services/pipeline.py --repo https://github.com/user/repo
# 启动 API 服务器
uvicorn backend.main:app --reload
# 启动 frontend dashboard
cd frontend && npm start
```
示例 JSON 输出：
```
{
"file": "app/routes/login.py",
"language": "python",
"findings": [
{
"line": 42,
"type": "SQL Injection",
"cwe": "CWE-89",
"confidence": 0.93,
"sources": ["static_analysis", "ml_classifier", "llm_reasoning"],
"explanation": "User input is concatenated directly into a SQL query without parameterization.",
"suggested_fix": "Use parameterized queries via the database driver's placeholder syntax."
}
]
}
```
## 模型训练（微调）
只有第2层涉及实际训练。高层步骤：
1. 下载并预处理 Devign/Big-Vul/SARD 数据集 (`data/dataset_loader.py`)
2. 使用 CodeBERT tokenizer 对函数进行标记化
3. 使用 `ml_model/train.py` 微调 CodeBERT 以进行二分类（有漏洞 / 无漏洞）
4. 使用 Weights & Biases 追踪实验（loss、accuracy、F1）
5. 在预留的测试集上进行评估 (`ml_model/evaluate.py`)
6. 将每种语言的检查点保存在 `ml_model/checkpoints/` 下
针对每种第1级语言（Python、Java、C/C++）单独重复此操作，因为每种语言都需要其自己微调的检查点。
## 评估指标
该框架使用以下指标进行评估：
- **Precision, Recall, F1-score** — 按语言、按漏洞类别
- **误报率** — LLM 推理层过滤静态分析输出之前与之后对比
- **层级贡献分析** — 每层独立捕获了多少真阳性
- **延迟** — 每个文件/repo 大小的平均扫描时间
结果记录在 `docs/evaluation_results.md` 中。
## API参考
| 端点 | 方法 | 描述 |
|---|---|---|
| `/scan/file` | POST | 上传单个文件进行扫描 |
| `/scan/repo` | POST | 提交 GitHub repo URL 进行扫描 |
| `/report/{scan_id}` | GET | 检索已完成扫描的发现结果 |
| `/report/{scan_id}/pdf` | GET | 下载 PDF 报告 |
| `/history` | GET | 列出过去的扫描 |
## 项目时间线
| 阶段 | 周数 | 交付物 |
|---|---|---|
| 文献回顾 + 数据集选择 | 1–2 | 确定数据集选择、repo 框架 |
| 静态分析层 | 3–4 | 可工作的 Semgrep/Bandit 集成 |
| ML分类器微调 | 5–8 | 训练好的 CodeBERT 检查点（第1级语言） |
| LLM推理层 | 9–10 | Prompt 设计、结构化输出解析 |
| 聚合器 + 评分 | 11–12 | 合并、加权的发现 pipeline |
| 后端 + 前端 | 13–14 | 可工作的 API 和 dashboard |
| 评估 + 撰写报告 | 15+ | 指标、最终报告、文档 |
## 局限性与未来工作
- ML 分类器的覆盖范围仅限于具有充足标记数据的语言；扩展到更多语言需要新的标记数据集或迁移学习方法。
- LLM 推理层依赖于外部 API —— 在大规模应用时会引入成本和延迟方面的考虑。
- 目前的设计分析的是函数级别的上下文；整个程序/跨文件的数据流漏洞（例如：跨文件的污点追踪）是未来可能的扩展方向。
- 未来的工作可以探索微调单一的多语言分类器（例如，使用 UniXcoder 或 CodeT5+），以减少对特定语言检查点的需求。
## 参考
- Zhou et al., "Devign: Effective Vulnerability Identification by Learning Comprehensive Program Semantics via Graph Neural Networks"
- Fan et al., "A C/C++ Code Vulnerability Dataset with Code Changes and CVE Summaries" (Big-Vul)
- Feng et al., "CodeBERT: A Pre-Trained Model for Programming and Natural Languages"
- Guo et al., "GraphCodeBERT: Pre-training Code Representations with Data Flow"
- Semgrep documentation: https://semgrep.dev/docs/
- NIST SARD / Juliet Test Suite: https://samate.nist.gov/SARD/
标签：DLL 劫持, 人工智能安全, 代码安全审计, 凭据扫描, 合规性, 大语言模型, 机器学习分类器, 测试用例, 请求拦截, 逆向工具, 错误基检测, 静态代码分析