为Smart Contract提取以AST、CFG、DFG为基础的异构图，使用图神经网络进行多粒度的分类，实现行级别以及合约级别的漏洞检测任务。
总体框架图如下所示：
得到的异构图案例如下所示：
在dataset.zip中已经压缩了28组数据集。分别是7种漏洞类型，每组漏洞类型都分为原始合约级别，增强以后合约级别、原始行级别、增强后行级别。每个文件夹中都存有对应的cmd命令，可用于运行、测试。
根据虚拟环境的复现过程操作以后，使用对应的虚拟环境以及参数调用对应的main.py即可。
(lunikhod) root@0efc54cd086a:~/MVD-HG# python main.py --help
usage: main.py [-h] [--run_mode RUN_MODE] [--create_corpus_mode CREATE_CORPUS_MODE] [--train_mode TRAIN_MODE] [--data_dir_name DATA_DIR_NAME] [--attack_type_name ATTACK_TYPE_NAME]
[--create_code_snippet] [--data_augmentation] [--gpu_id GPU_ID] [--coefficient COEFFICIENT] [--target_dir TARGET_DIR] [--target_file TARGET_FILE]
参数表
optional arguments:
-h, --help show this help message and exit
--run_mode RUN_MODE 运行模式: 1.create:创建数据集的时候用的。 2.train:训练模式(这个是文件级别的)。 3.predict:预测模式。 4.truncated:语言模型受损，需要重新训练。 5.line_classification_train:行级别的漏洞检测。
6.contract_classification_train:合约级别的漏洞检测。
--create_corpus_mode CREATE_CORPUS_MODE
创建文件的模式: 1.create_corpus_txt:仅仅创建语料库文件。 2.generate_all: 生成所有的向量文件。
--train_mode TRAIN_MODE
针对创建文件的模式一起使用的，如果创建文件的模式是generate_all那么这里应该标注是合约级别还是行级别，因为built_vector_dataset的时候需要根据不同的任务类型，读取不一样的json文件
--data_dir_name DATA_DIR_NAME
数据文件夹的名字，为了可以多进程启动运行:
--attack_type_name ATTACK_TYPE_NAME
本次要操作的漏洞的类型，专门只操作这种漏洞。
--create_code_snippet
本次操作是否是用来创建合约的片段的
--data_augmentation 本次操作是否是用来拓展数据集的
--gpu_id GPU_ID 本次操作的gpu用哪个
--coefficient COEFFICIENT
系数
--target_dir TARGET_DIR
目标文件夹
--target_file TARGET_FILE
目标文件数据集组成
原始合约级别漏洞检测结果，'\'代表不具备该漏洞类型的检测能力。
原始行级别漏洞检测结果，'\'代表不具备该漏洞类型的检测能力。
使用数据增强以后的合约级别漏洞检测结果，'\'代表不具备该漏洞类型的检测能力。
使用数据增强以后的行级别漏洞检测结果，'\'代表不具备该漏洞类型的检测能力。
不同类型组成异构图的消融实验结果(增强以后的数据集)。
进行参数实验以后得到的结果。
root@iZbp1dlnrh61z7i0vvirdpZ:~/code/MVD-HG# tree
.
├── append_control_flow_information.py # 为提取出的节点增加控制流信息。
├── append_data_flow_information.py # 为提取出的节点增加数据流信息。
├── append_method_message_by_dict.py # 为提取出的节点设置函数信息、所属Contract等。
├── bash # 一些实验时快速启动的脚本文件。
│ ├── go.sh
│ ├── 启动三次运行.vbs
│ ├── 启动虚拟环境.vbs
│ ├── 监视程序运行情况.vbs
│ ├── 直接启动tensorboard.vbs
│ └── 脚本说明.txt
├── bean
│ └── Node.py # 节点类。
├── built_corpus.py # 构建语料库文件的txt和pkl模型文件。
├── built_vector_dataset.py # 根据节点信息构建出持久化的本地向量化文件
├── classification_of_documents.py # 将所有有漏洞的文件进行分门别类地存放起来，方便下次检验。
├── compile_files.py # 编译智能合约源代码。
├── config.py # 配置文件。
├── contract_classification # 合约级别漏洞检测走的流程分支。
│ ├── contract_classification_dataset.py # 加载的数据集格式。
│ ├── contract_classification_model.py # 使用的模型。
│ └── contract_classification_train.py # 训练过程。
├── create_code_snippet.py
├── data # 存放实验时使用的数据源，以及各种生成的结果文件等。
│ ├── AST_json # 编译后的结果，保存json文件。
│ │ └── contract1_0xfD904a11fEC111F353ec8A5C9af203c59391dECA_Bank
│ │ └── 0xfD904a11fEC111F353ec8A5C9af203c59391dECA_Bank.sol
│ ├── attack # 之前提到的有漏洞的文件进行分门别类保存的地方。
│ │ ├── arithmetic_attack
│ │ │ └── contract1_0xb725213a735ae34e1903d1971700dd7f0858f212_BankofIsreal
│ │ │ └── 0xb725213a735ae34e1903d1971700dd7f0858f212_BankofIsreal.json
│ │ ├── dangerous_delegate_call_attack
│ │ │ └── contract1_0xfef736cfa3b884669a4e0efd6a081250cce228e7_Bob
│ │ │ └── 0xfef736cfa3b884669a4e0efd6a081250cce228e7_Bob.json
│ │ ├── reentry_attack_fold
│ │ │ └── contract1_0x064d0c8d8100ba8c57a63d75a5fb8ede18d7fe4b_QSHUCOIN
│ │ │ └── 0x064d0c8d8100ba8c57a63d75a5fb8ede18d7fe4b_QSHUCOIN.json
│ │ └── timestamp_attack_fold
│ │ └── contract1_0xF9d3402066E3a483f4ca7abFa78DEC61635E561f_PreCrowdsale
│ │ └── 0xF9d3402066E3a483f4ca7abFa78DEC61635E561f_PreCrowdsale.json
│ ├── complete # 当还是create_corpus_txt的时候，这里保存所有被处理过的文件，但是之后会被移送回sol_source。generate_all则不会移送回去。
│ │ └── contract1_0xfc9ec868f4c8c586d1bb7586870908cca53d5f38_KittyItemMarket
│ │ └── 0xfc9ec868f4c8c586d1bb7586870908cca53d5f38_KittyItemMarket.sol
│ ├── corpus_model.pkl # 语言模型。
│ ├── corpus.txt # 语言库。
│ ├── error # 如果编译出错，或者执行流程中出错的文件。
│ │ └── contract1_0xc9Fa8308cd98A6144450D68B6C546062dbBD984e_CrowdsaleTokenExt
│ │ └── 0xc9Fa8308cd98A6144450D68B6C546062dbBD984e_CrowdsaleTokenExt.sol
│ ├── idx_to_label.json # 通过自己的漏洞模型，给出的漏洞标签。
│ ├── img # 保存图形化以后的结果。
│ │ └── contract1_0xc9Fa8308cd98A6144450D68B6C546062dbBD984e_CrowdsaleTokenExt
│ │ └── 0xc9Fa8308cd98A6144450D68B6C546062dbBD984e_CrowdsaleTokenExt.svg
│ ├── log # 日志文件
│ ├── processed # 预处理以后形成的真正的数据集。
│ │ ├── graph_train.pt
│ │ ├── pre_filter.pt
│ │ └── pre_transform.pt
│ ├── raw # 原始的向量文件。
│ │ └── contract1_0xb725213a735ae34e1903d1971700dd7f0858f212_BankofIsreal # 每一个文件夹下面都会有三份边文件和一个节点文件。
│ │ ├── 0xb725213a735ae34e1903d1971700dd7f0858f212_BankofIsreal_ast_edge.json
│ │ ├── 0xb725213a735ae34e1903d1971700dd7f0858f212_BankofIsreal_cfg_edge.json
│ │ ├── 0xb725213a735ae34e1903d1971700dd7f0858f212_BankofIsreal_dfg_edge.json
│ │ └── 0xb725213a735ae34e1903d1971700dd7f0858f212_BankofIsreal_node.json
│ ├── sol_source # 原始文件夹，要处理的文件都放在这里面。
│ │ └── contract1_0xfc9ec868f4c8c586d1bb7586870908cca53d5f38_KittyItemMarket
│ │ └── 0xfc9ec868f4c8c586d1bb7586870908cca53d5f38_KittyItemMarket.sol
│ └── tensorboard_logs # 使用tensorboard以后保存的文件夹。
│ └── events.out.tfevents.1658990929.dell-PowerEdge-T640.126902.0
├── data_process.py # 额外的数据处理文件。
├── dataset.py # 构建数据集的文件。
├── dataset.zip # 实验使用的数据集。
├── hash_code.py # 可以过滤重复的文件，提取每个文件的hash结果。
├── line_classification # 行级别漏洞检测的流程分支。
│ ├── line_classification_dataset.py # 加载的数据集格式。
│ ├── line_classification_model.py # 使用的模型。
│ └── line_classification_train.py # 训练过程。
├── load_model_to_predict.py # 加载模型用于训练。
├── main.py # 主文件。
├── make_arithmetic_attack_label.py # 使用专家规则，求出arithmetic的标签文件（已弃用）。
├── make_reentry_attack_label.py # 使用专家规则，求出reentry的标签文件（已弃用）。
├── make_timestamp_attack_label.py # 使用专家规则，求出timestamp的标签文件（已弃用）。
├── merge.py
├── metric.py # 度量文件，在这里会计算多种的度量标准。
├── model.py # 网络模型（已弃用）。
├── print_tree.py # 将异构图按照树形打印出来，可视化出来。
├── read_compile.py # 读取编译完的结果，然后生成图结构的内存数据。
├── README.md
├── remove_blank.py # 删除源文件中的空行。
├── remove_comments.py # 删除源文件中的注释。
├── train.py # 训练文件（已弃用）。
└── utils.py # 一些通用的工具函数。徐敬杰
浙江工业大学 软件工程专业硕士在读
邮箱地址:925791559@qq.com