VulTriage

一种基于大语言模型和静态代码分析的智能漏洞检测框架，支持多数据集评估和消融研究。

📋 目录

• 介绍
• 主要特性
• 系统要求
• 快速开始
• 使用方法
• 数据集与模型
• 常见问题

介绍

本项目是一个基于GPT-4o和BGE-M3模型的智能漏洞检测框架，结合了：

• 静态结构分析：控制流图（CFG）、数据流图（DFG）、抽象语法树（AST）
• 检索增强生成（RAG）：基于CWE知识库的漏洞模式匹配
• 代码理解：使用大语言模型进行语义分析
• 成对/单样本检测：支持多种评估模式

主要特性

• ✅ 支持PrimeVul和Kotlin数据集
• ✅ 集成CWE漏洞知识库检索
• ✅ 可配置消融研究
• ✅ 实时指标追踪
• ✅ 灵活的超参数实验支持

系统要求

• Python >= 3.9
• CUDA（可选，用于GPU加速）

快速开始

1. 安装依赖

pip install -r requirements.txt

2. 准备资源

下载BGE-M3模型

模型会自动下载，你也可以手动下载：

git clone https://huggingface.co/BAAI/bge-m3
# 注意：BGE-M3相关依赖可通过FlagEmbedding安装，详情请参考：pip install -U FlagEmbedding

准备CWE数据

确保项目根目录下存在 cwec_latest.xml/cwec_v4.19.1.xml 和 1435.csv/1435_filtered_columns.csv（CWE索引数据）文件。

注意：本项目使用CWE Top 25最危险软件弱点。相关文件已包含在项目根目录中。

准备数据集

将以下文件放入项目根目录：

• primevul_test_paired.jsonl（PrimeVul数据集）
• LLM4Vul-main/LLM4Vul-main/data/kotlin_data.csv（Kotlin数据集）

3. 配置API密钥

在 main.py 或 hyperparam_exp.py 中修改OpenAI API配置：

client = OpenAI(
    api_key="YOUR_API_KEY",
    base_url="https://api.poixe.com/v1"
)

注意：当前URL https://api.poixe.com/v1 解析失败。可能是无法支持的网页类型或链接已断开。请验证链接有效性或切换到其他API提供商。

4. 运行漏洞检测

基础模式（完整功能）

python main.py --wotask all --jsonl_file primevul_test_paired.jsonl

消融研究模式

# 不含控制流信息
python main.py --wotask wo_ctrl_info_pth --jsonl_file primevul_test_paired.jsonl
# 不含RAG知识
python main.py --wotask wo_rag_pth --jsonl_file primevul_test_paired.jsonl
# 不含代码理解
python main.py --wotask wo_exp_pth --jsonl_file primevul_test_paired.jsonl
# 不含任何辅助信息
python main.py --wotask nan --jsonl_file primevul_test_paired.jsonl

超参数实验

# 成对评估模式
python hyperparam_exp.py --mode pairwise --input_file primevul_test_paired.jsonl --detail_level C --info_comb all --desc_level concise --max_vuln 2
# 单样本评估模式
python hyperparam_exp.py --mode single --input_file LLM4Vul-main/LLM4Vul-main/data/kotlin_data.csv

5. 查看结果

检测结果实时保存到 metrics_log.txt，最终评估摘要输出到终端：

• 评估指标包括准确率、召回率、假阳性率等
• 消融研究结果比较不同模式下的性能差异，以分析各模块的贡献

使用方法

主要脚本概览

参数说明

main.py 参数

hyperparam_exp.py 参数

数据集与模型

📊 数据集

1. PrimeVul数据集

• 引用：Yangruibo Ding, Yanjun Fu, Omniyyah Ibrahim, Chawin Sitawarin, Xinyun Chen, Basel Alomair, David Wagner, Baishakhi Ray, 和 Yizheng Chen. 2025. 使用代码语言模型检测漏洞：我们走了多远？ISSTA 2025 (2025).
• 用途：成对漏洞检测基准
• 格式：JSONL（包含 func 和 target 字段）
• 来源：ISSTA 2025 会议（ISSTA官方网站）。有关数据集访问权限，请参考原始论文。

2. Kotlin数据集

• 引用：Triet Huynh Minh Le, M Ali Babar, 和 Tung Hoang Thai. 2024. 低资源语言中的软件漏洞预测：CodeBERT和ChatGPT的实证研究. 在《第28届国际软件工程评估与评价会议论文集》中，第679–685页.
• 用途：低资源语言漏洞预测研究
• 格式：CSV（包含 code 和 label 列）
• 来源：EASE 2024 会议（EASE 2024官方网站）。有关数据集访问权限，请参考原始论文。

3. CWE数据库

• 版本：CWE v4.19.1
• 来源：MITRE CWE
• 用途：漏洞模式知识库和检索增强
• 注意：CWE可自由用于研究、开发和商业用途。必须保留对MITRE的版权归属。

🤖 模型

1. BGE-M3嵌入模型

• 开发者：BAAI（北京人工智能研究院）
• HuggingFace：BAAI/bge-m3
• 许可证：MIT许可证（注意：当前模型许可证页面解析失败。请参考FlagEmbedding项目许可证，为MIT许可证）
• 用途：稠密和稀疏向量混合检索，支持多语言、多粒度（最大输入长度8192个token）和多功能检索
• 相关项目：FlagEmbedding（BGE系列模型的开源仓库）

2. GPT-4o

• 开发者：OpenAI
• 模型版本：gpt-4o-2024-11-20
• 许可证：遵循OpenAI使用条款（自2026年1月1日起生效）
• 用途：代码理解、漏洞识别、推理与判断
• 注意：使用须遵守OpenAI服务条款。禁止非法或有害活动以及模型逆向工程。

📦 依赖库

常见问题

Q1：如何更换OpenAI API提供商？

修改 main.py 或 hyperparam_exp.py 中的 base_url 和 api_key。

Q2：为什么检索结果不准确？

1. 检查 1435.csv/1435_filtered_columns.csv 是否正确加载，确保CWE索引数据格式正确；
2. 确认BGE-M3模型已正确下载。可以通过 git clone https://huggingface.co/BAAI/bge-m3 手动下载。

Q3：如何自定义评估指标？

修改 metrics_tracker.py 中的 update_pair_metrics 方法，添加或调整指标计算逻辑。

Q4：如果BGE-M3模型下载失败怎么办？

直接从HuggingFace网站下载模型文件，解压到项目根目录，或检查网络连接后重新执行 git clone https://huggingface.co/BAAI/bge-m3 命令。