图神经网络驱动的切片级漏洞检测：从PDG到GNNExplainer

简介：这份源码与项目说明包面向软件安全方向的毕业设计、课程设计及期末大作业，聚焦基于图神经网络的切片级漏洞检测与解释任务。包内共263个文件，以90个Python脚本为主，辅以pyc字节码、zbak备份、JSON与DOT图结构文件等，整体约5MB，涵盖数据预处理、代码标准化、Joern生成PDG、代码切片构建、w2v向量嵌入、漏洞检测模型训练以及GNNExplainer/PGExplainer解释改进等完整流程。项目中提供了可直接运行的源码、路径配置说明、漏洞实例的源代码及解释结果行号，并附带常见问题处理脚本，便于读者复现实验并开展二次开发。目前已有80人学习浏览，适合具备Python与深度学习基础、希望快速搭建漏洞检测原型的本科生和研究生参考。资源来源于网络分享，仅供学习交流使用。

1. 图神经网络驱动的切片级漏洞检测，把安全分析从逐行读代码变成看依赖图

传统漏洞检测工具多基于正则或AST模式匹配，遇到指针别名、跨函数数据流时基本失效。切片级漏洞检测另辟蹊径：把一个函数或代码片段中与敏感操作（如memcpy、strcpy）有依赖关系的行组织成程序依赖图（PDG），再用图神经网络学习缺陷模式。这套基于Python的实现，从NVD漏洞数据出发，经过Joern生成PDG、W2V嵌入、GNN分类、GNNExplainer解释，形成了一条完整闭环。如果你正在做毕业设计、课程设计或者期末大作业，想找一份“能跑通、有解释、还能扩展”的图神经网络源码，这份项目说明值得逐行拆。你不需要先精通编译器，只要会按命令操作Python脚本，并理解数据流走向即可。下面我按数据流动顺序，讲清楚每个环节为什么那样设计、参数怎么调、失败时看哪里。

2. 数据预处理：从NVD漏洞行到Joern PDG图

2.1 原始数据格式与漏洞行提取

项目开始先读取一个CSV文件，每个样本包含CVE编号、函数名、修复前后代码等信息。raw_data_preprocess.py负责提取“漏洞减号行”——也就是补丁中带-的删除行，这些行往往就是缺陷所在。程序会把提取结果整理成nvd_vul_lineinfo.json，格式大致是：

{ "CVE-2019-1234": { "file": "xen/arch/x86/foo.c", "func_name": "bar", "vul_lines": [10, 12] } }

file是源文件路径，func_name是漏洞函数名，vul_lines是漏洞行号列表。这个JSON是后续所有处理的锚点：切片要以它为终点，评估要以它为真值。如果CSV里一个函数有多个减号行，就全部记入vul_lines，但同一个函数只保留一个图样本。

提示：vul_lines必须与标准化后的代码行号一一对应。一旦预处理阶段改了行号，这里要做偏移修正，否则后面切片会定位到错误的节点。

2.2 代码标准化：消除命名噪音

漏洞模式往往存在多种等价的命名写法，比如size、len、nbytes其实都是长度变量。直接拿原始代码训练，GNN会把大量参数学在变量名上，泛化性很差。code_normalize/normalization.py的作用就是把变量名替换成var_0、函数名替换成func_0，同时保留关键字、运算符和字面量。这样模型学到的是结构缺陷，而不是某个项目的命名习惯。运行方式很简单：

python preprocess/code_normalize/normalization.py

脚本内部需要改源目录与目标目录两个变量。标准化后的代码会生成在指定输出目录，注意目录层级不要和原始目录混在一起，否则Joern会把多个版本解析进同一个CPG，导致节点ID错乱。

2.3 Joern三阶段生成PDG与LineInfo

Joern是生成代码属性图（CPG）的利器，能一次性产出AST、CFG、PDG。PDG融合了数据依赖边和控制依赖边，正是切片和GNN需要的拓扑结构。项目里写好了joern_graph_gen.py，分三个阶段执行。

首先解析代码，生成二进制中间结果：

python preprocess/joern_graph_gen.py -i cleaned_src -o joern_out -t parse

-i是标准化代码目录，-o是Joern工作目录，-t parse告诉脚本只做解析。解析完成后，joern_out里会有cpg.bin文件。接着导出PDG：

python preprocess/joern_graph_gen.py -i joern_out -o export_pdg -t export -r pdg

这里-i直接指向包含cpg.bin的目录，而不是源代码目录；-r pdg表示导出格式。每个函数会生成一个xxx.dot文件，文件名形如1_xen-4.12.0-4@API#4.dot，其中包含了函数内各语句节点的依赖边。

最后导出行号信息：

python preprocess/joern_graph_gen.py -i joern_out -o export_pdg -t export -r lineinfo_json

这步生成lineinfo.json，记录每个图节点对应的源码文件名、行号、列号、变量名和操作符。它是后续把解释结果映射回漏洞行的关键桥梁。三阶段的参数总结如下：

易错点在于第二次和第三次执行时，-i应该指向Joern工作目录而不是源码目录。如果你把-i写成源码路径，Joern会重新解析一遍，生成的节点ID会和lineinfo.json对不上。检查方法很直接：打开export_pdg下的任意.dot文件，看节点注释里是否出现正确的源码行号；再看lineinfo.json里是否包含同样的节点ID。

3. 代码切片与W2V嵌入：把PDG变成向量输入

3.1 以漏洞行为终点的后向切片

有了PDG和行号映射，下一步是生成代码切片。项目里的slice_preocess/main.py读取三个输入：lineinfo.json、PDG的.dot文件以及vul_lines字典，输出complete_pdg和slice_pdg两份结果。切片的基本思想是：从漏洞行对应的节点出发，沿依赖边后向遍历，收集所有影响漏洞点的语句。这里给出一个最小后向切片的伪代码实现：

def backward_slice(adj_rev, vul_node): # adj_rev: 逆向邻接表, 字典: 节点 -> 前驱节点列表 visited = set() stack = [vul_node] while stack: n = stack.pop() if n in visited: continue visited.add(n) for pred in adj_rev.get(n, []): stack.append(pred) return visited

adj_rev可以直接从.dot文件解析，把边方向反转；vul_node是通过lineinfo.json查到的漏洞行对应节点。返回值是所有可达前驱的节点集合，这就是后向切片。为什么只做后向？因为漏洞行为是被污点数据影响的，前向切片用于追踪影响传播，而判定漏洞根因时后向切片更紧凑，能减少无关节点的干扰。

切片之后，slice_pdg保留了节点属性和边关系，每个节点还带有从lineinfo.json继承的行号与操作符信息。这个阶段的产物已经是一张以漏洞行为焦点的小型程序依赖图。

3.2 训练Word2Vec生成节点特征

GNN不能直接吃代码文本，需要把每个节点表示成稠密向量。项目采用Word2Vec对切片中的token序列建模。train_w2v.py会读取所有切片，把每个节点的代码片段按标识符、关键字、运算符切分成token序列，训练一个CBOW或Skip-gram模型。典型配置如下：

python preprocess/train_w2v.py --embedding_dim 128 --window 5 --min_count 1

embedding_dim是最终节点向量维度，建议与后续GNN的hidden_dim保持一致；window是上下文窗口，取5能兼顾局部语法和短距离依赖；min_count设为1保证低频标识符不丢失，因为漏洞代码中的特殊变量名往往很关键。训练结束后，每个token得到一个向量，而节点向量由该节点包含的所有token向量求和或平均得到。

3.3 从PDG到Devign格式的图样本

得到节点向量后，joern_to_devign脚本负责把slice_pdg转换成图神经网络标准输入格式。它输出三个数组：特征矩阵x、邻接矩阵adj和标签y。之所以称为Devign格式，是因为这一格式来自图神经网络漏洞检测的经典工作Devign，很多开源模型都按这个格式加载数据。

一般我会在每个图样本里额外记录节点ID与行号的映射表，这样后面解释模型输出节点重要性时，可以直接对应到源码行。这一步看似简单，但如果你在切片后重新去掉了某些孤立节点，映射表也要同步更新，否则解释结果会偏移几行。

4. 漏洞检测模型训练：Devign架构与参数调整

4.1 为什么选门控图卷积作为编码器

项目使用的漏洞检测模型是典型的编码器-分类器结构。编码器采用门控图卷积（GGNN），它对每个节点通过多轮消息传递聚合邻居信息，并在传递过程中用GRU控制信息更新。相比普通GCN，GGNN能更好地建模依赖边的方向性——PDG里的数据依赖和控制依赖语义不同，门控机制可以自适应学习两种边的权重。分类器则是简单的全连接加Softmax，输出正常或漏洞二分类概率。

4.2 数据划分与模型入口

打开slice_level_model/main.py，要改三处：训练集/测试集文件名、数据目录、模型保存路径。数据划分建议按CVE级别进行，同一个CVE的多个切片必须全放同一边，否则模型会通过记忆CVE编号产生虚假的高准确率，这就是典型的数据泄漏。

python slice_level_model/main.py --train_path data/train.txt --test_path data/test.txt --epochs 100 --lr 0.001 --batch_size 64

如果你的实现里没有用argparse，直接把main.py顶部的TRAIN_FILE、TEST_FILE、CKPT_DIR改成实际路径即可。训练过程中每轮结束都会在测试集上计算准确率、F1，并保存最佳模型到ckpt目录。

4.3 训练超参速查表

深度学习训练有很强的随机性，项目说明也提到“训练的数据可能和源码作者有出入”。为了尽可能复现，固定随机种子并统一初始化方式很重要。我一般在main.py开头加：

import torch, numpy as np, random torch.manual_seed(42) np.random.seed(42) random.seed(42)

如果训练出的F1远低于预期，先检查数据划分是否泄漏，再检查邻接矩阵是否正确归一化，最后再看学习率是否过大导致损失爆炸。

5. GNNExplainer与PGExplainer：从预测结果反推漏洞行

5.1 解释目标：节点重要度掩码

漏洞检测模型只给出二分类结果是不够的，安全人员需要知道模型依据哪些代码判断为漏洞。GNNExplainer的做法是为每个输入图学习一个节点掩码和边掩码，目标是最大化原始预测概率与掩码后预测概率的互信息。掩码值越高，说明该节点对分类决策的贡献越大。项目在benchmar/kernal/pipeline.py中封装好了完整的解释流程，参数集中在args.py里。

5.2 运行pipeline.py与args.py配置

主要的加载部分有三步：数据集加载（带解释的切片，也就是完成了节点嵌入的图样本）、模型加载（选择第4章保存的ckpt文件）、输出路径设置。args.py里的关键项如下：

parser.add_argument("--data_dir", default="data/explain_slices", help="带嵌入的切片数据集") parser.add_argument("--model_ckpt", default="checkpoint/best_f1.pth", help="漏洞检测模型权重") parser.add_argument("--output_dir", default="results/important_nodes", help="解释结果保存目录") parser.add_argument("--explainer_type", default="gnn_explainer", help="gnn_explainer 或 pg_explainer")

设置好之后，运行：

python benchmar/kernal/pipeline.py

脚本会为每个测试样本输出一个importance_scores.json，里面记录每个节点ID及其重要性分数。注意这里的benchmar/kernal是项目给定的目录名，如果你从别的途径获得的源码里写的是benchmark/kernel，以实际仓库为准。

5.3 PGE替换与行号映射

PGExplainer与GNNExplainer的思路不同，它训练一个全局解释模型，一次性预测所有图的边重要性，速度更快，适合大规模图。替换方式极为简单，在args.py里把explainer_type改成pg_explainer，再保证数据路径和模型路径不变即可，其余代码无需改动。如果遇到解释结果为空，多半是图数据里没有边，检查切片后的.dot是否在迁移过程中丢失了边。

拿到重要性分数后，需要执行lineinfo_dict.py，它从lineinfo.json中读出节点ID到行号的映射。然后就可以把每个样本的top-k重要节点翻译成行号列表。interpre_example目录里提供了RQ3实例的完整材料：源代码、切片dot文件以及解释结果行号，非常适合对照检查你的流程是否跑对。

6. 排错与效果验证：dot文件修复与解释质量评估

6.1 dot文件加载错误与修复

Joern导出的.dot文件偶尔会包含格式不规范的节点ID，比如带引号或特殊字符，导致Python的pydot或networkx加载失败。项目提供了dot_fix.py，直接执行：

python preprocess/dot_fix.py --input bad.dot --output fixed.dot

它会统一转义节点名，去除重复边，并补全缺失的}。如果你的图数据规模较大，建议在批量处理后抽查几个文件，确保节点数量与lineinfo.json对得上。

6.2 解释效果量化：Top-K命中率

解释效果评定代码在intrepre_effect.py中，在此之前必须先执行lineinfo_dict.py生成行号字典。评估指标常用Top-K命中率：把模型解释出的节点按重要性排序，取前K个，计算这些节点对应的行号集合与真实漏洞行集合的交集比例。输出的表格大致如此：

上表只是结构示意，具体数值依赖你的训练数据和随机种子。如果命中率偏低，先看切片是否太宽泛——切片节点过多会稀释Top-K的准确率，适当缩小后向切片的深度或采用语句级去重会有效果。

6.3 使用Graphviz着色解释结果

最后分享一个实用技巧：把重要性分数映射到Graphviz的颜色深浅，直接渲染出可视化的漏洞子图。修改切片dot文件中的fillcolor属性，让解释出的重要节点标红，再运行dot -Tpng slice.dot -o output.png生成图片。这一步能直观看出模型依赖的是赋值语句还是条件判断，也能帮助排查是否出现解释到无关库函数的问题。在Graphviz渲染时，如果节点过多，先用tred对子图做传递约简，只保留关键依赖路径，不然图片会连成一片。更细的调试手段是打印解释top-1节点对应的源码行，人工复核该行是否真正参与数据依赖，这一步能快速发现切片或嵌入环节的隐性错误。

本文还有配套的精品资源，点击获取