AMPLE实验复现

  • A+
所属分类:笔记

记录实验复现流程,包括预处理步骤,遇到的错误及解决方案。

原文链接:Vulnerability Detection with Graph Simplification and Enhanced Graph Representation Learning | IEEE Conference Publication | IEEE Xplore

项目地址:https://github.com/AMPLE001/AMPLE

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

发表评论

:?: :razz: :sad: :evil: :!: :smile: :oops: :grin: :eek: :shock: :???: :cool: :lol: :mad: :twisted: :roll: :wink: :idea: :arrow: :neutral: :cry: :mrgreen: