SuZeAI/solidity-analyzer
GitHub: SuZeAI/solidity-analyzer
结合CodeBERT与图神经网络的智能合约漏洞检测流水线，用于识别重入、时间戳依赖等Solidity合约安全缺陷。
Stars: 0 | Forks: 0
# GraphFusionVulDetect：智能合约漏洞检测流水线（这是前端源代码）
本代码库包含了一个用于检测智能合约漏洞的流水线的实现及相关研究。该流水线整合了使用 **CodeBERT** 嵌入进行的函数级分析，以及使用 **图神经网络** 进行的基于图的分类。
## 项目描述
本项目的目标是开发一个有效且可扩展的系统，用于检测智能合约中的漏洞。该过程包括：
1. **函数级漏洞分类**：
- 分析智能合约中的函数以识别特定的漏洞类型（例如，重入、时间戳依赖、未处理异常）。
- 创建一个数据集，将函数标记为易受攻击或不易受攻击。
- 针对每种漏洞类型对 **CodeBERT** 进行微调，以生成代表功能语义的嵌入。
2. **基于图的分类**：
- 分析源代码以构建 **函数调用图**，其中：
- 节点代表函数。
- 边代表函数调用关系。
- 使用微调后的 **CodeBERT** 模型生成节点嵌入。
- 这些图作为 **GNN 模型**（例如，GCN、GraphSAGE、GAT）的输入，用于将智能合约分类为易受攻击或不易受攻击。
3. **基线对比**：
- 将结果与现有方法进行比较，例如 CBGRU、Peculiar、VulBERTa、TMP、AME、MANDO、MANDO-HGT 等，以展示所提出方法的性能提升。
## 关键特性
- **数据收集**：
- 使用了诸如 `smartbugs/smartbugs-wild` 和 `mwritescode/slither-audited-smart-contracts` 等数据集。
- 使用 **Slither**、**SmartCheck** 和 **Mythril** 等工具对函数和源代码进行预处理。
- **建模技术**：
- 微调 **CodeBERT** 以生成嵌入。
- **GNN** 用于基于图的分类。
- 与替代架构进行比较，包括混合模型（例如，CBGRU）。
- **优化**：
- 处理输入 token 长度的限制（例如，CodeBERT 的 512 token 限制）。
- 研究了诸如选择关键 token（前 128 + 后 382 个 token）等策略。
## 结果
- 在检测漏洞方面观察到显著改进，特别是在重入漏洞方面，优于传统模型。
- 对不同的图构建技术和特征增强进行了实验。
## 文件夹结构
- `datasets/`：包含用于函数和图的数据集。
- `core/`：流水线组件的核心实现，包括：
- `finetune_llm/`：CodeBERT 微调、数据集、流水线和组件。
- `GFD/`：GFD 模型流水线、数据处理和模块。
- `preprocessing/`：用于数据集预处理和分析的脚本。
- `utils/`：流水线的工具函数。
- `notebooks/`：用于数据处理、模型训练和评估的 Jupyter notebooks。
- `docs/`：额外的项目文档和基线对比。
## 开始使用
1. **前置条件**：
- Python 3.8+
- PyTorch
- Hugging Face Transformers
- DGL (Deep Graph Library)
- Pandas, NumPy, Scikit-learn
2. **设置**：
克隆代码库并安装依赖项：
git clone https://github.com/QuangNguyen2910/GraphFusionVulDetect.git
cd GraphFusionVulDetect
pip install -r requirements.txt
3. **数据准备**：
- 按照 `docs/data_preparation.md` 中的说明预处理数据集并创建 FCG。
4. **训练与评估（使用 Makefile）**：
你可以使用提供的 Makefile 轻松运行训练和评估命令：
- 微调 CodeBERT：
make finetune-train
- 评估 CodeBERT：
make finetune-eval
- 训练 GFD 模型：
make gfd-train
- 评估 GFD 模型：
make gfd-eval
这些命令将执行相应的 Python 模块以进行训练和评估。
## 引用
如果你在研究中使用了本项目，请引用配套论文：
```
@article{GraphFusionVulDetect,
title={GraphFusionVulDetect: A Smart Contract Vulnerability Detection Using CodeBERT and GNNs},
author={Quang Nguyen, Tuyen Vu, Minh Pham, Kien Nguyen, Cong Tran},
journal={None},
year={2024}
}
```
## 联系方式
如有问题或合作意向，请联系 [Quang Nguyen](nguyenquang71103@gmail.com)。
标签：Apex, CodeBERT, FCG, GAT, GCN, GNN, GraphSAGE, Mutation, Mythril, Slither, Solidity, 云安全监控, 代码嵌入, 凭据扫描, 函数调用图, 前端, 区块链安全, 可视化, 图神经网络, 时间戳依赖, 智能合约, 机器学习, 深度学习, 网络安全, 自动化攻击, 逆向工具, 重入攻击, 隐私保护, 静态分析