- A+
记录实验复现流程,包括预处理步骤,遇到的错误及解决方案。
1 预处理
1.1 joern工具下载
首先从作者提供的网址:https://zenodo.org/records/7323504/files/joern.zip?download=1 下载编译好的joern 工具。(不同版本的 joern 生成的 cpg 图可能存在一定的差异,为尽可能复现论文中报告的实验结果,建议使用作者提供的 joern 版本)。
解压至 data_processing 文件夹中,此时的文件夹结构应为:
data_processing
├─ cpg_GS.py
├─ cpg_original.py
├─ cpg_TGS.py
├─ cpg_VGS.py
├─ joern.md
├─ process.py
├─ utils.py
├─ word2vec.py
└─ joern
├─ process.py
├─ readme.txt
├─ slicer.sh
└─ joern
└─ joern
├─ .gitignore
├─ AUTHORS
├─ build.gradle
├─ build.sh
├─ joern-parse
├─ ....
├─ ....
1.2 数据集预处理
下载 Devign 数据集:function.json ,(需科学上网,文中采用了 Devign, Reveal 和 BigVul 三种数据集,本文以 Devign 数据集为例进行演示),并放在AMPLE 根目录中。
新建 dataset_preprocess.py 文件,用于提取 function.json 中的每一个函数样本,并按照 0.c, 1.c, 2.c, ... 的命名方式生成 C 文件:
import json
import os
# 读取 function.json 文件
with open('function.json', 'r', encoding='utf-8') as file:
data = json.load(file)
# 提取所有键为 "func" 的值
func_values = []
# 递归提取所有键为 "func" 的值
def extract_func_values(data):
if isinstance(data, dict): # 如果是字典
for key, value in data.items():
if key == "func":
func_values.append(value)
else:
extract_func_values(value) # 递归处理嵌套的字典
elif isinstance(data, list): # 如果是列表
for item in data:
extract_func_values(item) # 递归处理列表中的每个元素
# 调用提取函数
extract_func_values(data)
# 创建 devign_dataset 文件夹,如果不存在的话
output_folder = 'devign_dataset/devign_raw_code'
os.makedirs(output_folder, exist_ok=True)
# 为每个提取的值生成一个 .c 文件并保存到 devign_dataset 文件夹
for i, func_value in enumerate(func_values):
filename = os.path.join(output_folder, f"{i}.c")
# 将 func_value 写入 .c 文件
with open(filename, 'w', encoding='utf-8') as f:
f.write(func_value)
print(f"已生成文件: {filename}")
1.3 代码路径修改
原代码中存在一些路径问题,具体修改为:
# AMPLE/data_processing/process.py 文件:
# 第 24 行
"bash ./slicer.sh " 修改为 "bash ./joern/slicer.sh " # 注意后面的空格
# 第 38 行
"../devign_dataset" 修改为 "../devign_dataset/devign_raw_code"
# AMPLE/data_processing/joern/slicer.sh 文件,第 19 行
./joern/joern-parse 修改为 ./joern/joern/joern/joern-parse
执行下列命令,开始调用 joern 解析源代码(解析过程相当漫长,大约4-5小时,慢慢等待即可...):
cd data_processing
python process.py --start 0 --end 27317
如果执行上述命令遇到 ./joern/slicer.sh: line 19: ./joern/joern/joern/joern-parse: Permission denied 等错误,通过下列命令修改执行权限后再次运行即可:
chmod +x ./joern/joern/joern/joern-parse
2 图简化
安装作者提供的库版本,避免因 torchdata 和 dgl 库版本不符导致的报错:
conda install torch==1.9.0
conda install pytorch::torchdata
conda install pyyaml
conda install pydantic
下载作者训练好的 word2vec 模型:https://zenodo.org/records/7333062#.Y3c5SHZByUk,放入文件夹`devign_dataset/devign_wv_models`中。
新建`devign_dataset/devign_parsed`文件夹,并将`data_processing`文件夹下生成的每个解析结果移动到该文件夹中。
将devign数据集按照8:1:1的比例划分为train、valid、test三个部分,并放入文件夹`devign_dataset/devign_data_split`中。
修改 /AMPLE_code/data_loader/batch_graph.py
以及 /AMPLE_code/data_loader/dataset.py
中的DGLGraph()
为 graph()
,以消除下列报警问题:
/home/lxh/anaconda3/envs/ample/lib/python3.9/site-packages/dgl/heterograph.py:72: DGLWarning: Recommend creating graphs by dgl.graph(data) instead of dgl.DGLGraph(data).dgl_warning('Recommend creating graphs by dgl.graph(data)'
# 修改 661、664、668 为按行读取json
# 修改第 649-653 行的超参数路径
# 修改 702 行为 file_name = str(i) + '.c'
# 修改 705 行为 label = int(entry['target'])
python cpg_GS.py
执行完毕后会在`devign_dataset/decign_cpg_c2_2`文件夹中生成9个文件:
bad_file.json
devign-test-num-v2.json
devign-test-v2.json
devign-train-file.json
devign-train-num-v2.json
devign-train-v2.json
devign-valid-file.json
devign-valid-num-v2.json
devign-valid-v2.json
3 模型训练
安装torch(建议安装此版本,避免dgl冲突问题):
pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio==0.9.0 -f https://download.pytorch.org/whl/cu111/torch_stable.html
安装dgl-gpu版本(注意torch和cuda版本对应):
pip install dgl-cu111 -f https://data.dgl.ai/wheels/repo.html
开始训练模型:
cd AMPLE_code
python main.py