溯源图与图神经网络：APT攻击检测的RGAT+GRU实战解析

简介：这是一份2023年华中科技大学毕业设计项目，聚焦基于溯源图的APT攻击检测方法优化，适合计算机相关专业学生完成毕设、课设或科研进阶。项目以Python为主要开发语言，使用DARPA TC5 Cadets与StreamSpot数据集，包含RGAT、GRU等模型实现与对比，覆盖数据预处理、模型训练、效果评估等完整流程。压缩包共25个文件，包括11个Python脚本、7个XML工程配置、4个Markdown说明文档，另含模型文件与忽略配置，整体约45KB，目录结构清晰、便于按模块阅读。已有141人浏览学习。通过该资源可获取可直接运行的源码、数据集处理思路、模型优化细节及配套文档说明，适合希望在APT检测方向快速上手并做扩展完善的学习者。

1. 溯源图与 APT 攻击检测：这个毕设项目到底给了你什么

APT（高级持续性威胁）攻击最让人头疼的不是某个漏洞有多隐蔽，而是攻击者进入内网后可以在你眼皮底下潜伏几个月，一步一步横向移动、提权、窃取数据。传统的特征检测对这种慢速、低漏洞的攻击基本失效，所以最近几年安全研究者都把目光转向了溯源图（Provenance Graph）：把系统里的进程、文件、网络连接全部变成节点，把系统调用变成边，这样一次攻击就算拆成再小的动作，也会在图上留下一条可以回溯的路径。2023年华中科技大学这个毕业设计做的正是这件事，用 RGAT 加 GRU 对溯源图做建模，在 DARPA TC5 Cadets 和 StreamSpot 两个公开数据集上跑通并优化了 APT 检测流程。对想做安全方向毕设、或者想把图神经网络用到异常检测里的同学来说，这份源码最大的价值不是模型有多新，而是一套完整可复现的“数据预处理 → 构图 → 模型训练 → 评估”流程，拿到手能直接改、直接跑。

2. 从审计日志到溯源图：数据预处理与图构建

2.1 两个数据集：DARPA TC5 Cadets 和 StreamSpot 的差异

要复现这个项目，第一件事就是搞清楚它喂给模型的数据长什么样。项目里用了两个数据集：DARPA TC5 Cadets 和 StreamSpot。DARPA TC5 是 DARPA 透明计算项目发布的系统审计数据集，Cadets 是其中一个场景，记录的是 Linux 系统上的进程执行、文件读写、网络连接等系统调用级事件。StreamSpot 是另一个安全检测常用数据集，内容是浏览器、下载器等应用的系统调用流，更偏向流式检测场景。

这两个数据集各有各的脾气。DARPA TC5 Cadets 的数据量比较大，raw 格式是 json，每条记录包含进程 PID、PPID、操作类型、文件路径、时间戳等字段；StreamSpot 的数据则比较规整，是图格式的，节点和边都预先标注好。所以项目里你会发现代码分成了darpa_cadets_RGAT.py和streamspot_RGAT.py两套处理逻辑，原因就是两个数据集的解析方式完全不同。

在动手跑之前，建议先分别打开这两个文件看一眼，把数据集路径改成你自己的绝对路径。我第一次跑的时候直接把 StreamSpot 的解析逻辑套到了 DARPA 上，结果边列表全是空的，模型训练出来的 loss 直接不下降。后来才意识到两个数据集的事件字段名称都不一样，必须分开处理。

2.2 图构建的核心代码：节点、边与时间窗口

不管哪个数据集，最终都要转成一张有向图。节点是进程、文件、socket 等实体，边是它们之间的调用关系，比如fork、execve、write。实际代码里一般会先读一条条审计事件，再聚合成边列表。下面这段是典型的 DARPA 日志解析逻辑，项目里darpa_cadets_RGAT.py的预处理部分就是这个思路：

import json from collections import defaultdict def build_graph_from_darpa(log_file, time_window=60): """ 从 DARPA TC5 的 JSON 审计日志构建溯源图 :param log_file: 审计日志路径 :param time_window: 时间窗口（秒），用于切分子图 :return: node_list, edge_list, timestamp_list """ edges = [] nodes = set() time_index = [] with open(log_file, 'r', encoding='utf-8') as f: for line in f: record = json.loads(line.strip()) # 关键字段：进程对象、文件对象、操作类型 subject = record.get('subject', {}) # 主体：进程 obj = record.get('object', {}) # 客体：文件/套接字 op = record.get('op', 'unknown') # 操作类型 ts = record.get('timestamp', 0) src = subject.get('pid', 'unknown') dst = obj.get('path') or obj.get('fid') or obj.get('socket') if not dst: continue nodes.add(src) nodes.add(dst) edges.append((src, dst, op, ts)) time_index.append(ts) # 按时间窗口切分成多个子图，每个子图是一个快照 snapshot_edges = [] current_window_start = time_index[0] current_edges = [] for e in edges: if e[3] - current_window_start <= time_window: current_edges.append(e) else: snapshot_edges.append(current_edges) current_edges = [e] current_window_start = e[3] if current_edges: snapshot_edges.append(current_edges) return list(nodes), snapshot_edges

这段代码的关键在于time_window参数。APT 攻击的路径可能跨越很长时间，窗口太短会把一条完整攻击链切成碎片；窗口太长又会把无关进程混进来，导致图过大、内存爆炸。我一般习惯先用一小段日志跑一遍，统计事件时间戳的跨度，再取中间值作为窗口大小。项目里默认是 60 秒，StreamSpot 数据集本身已经是图结构，不需要这么切，直接读预建的邻接表就行。

3. 模型方案：RGAT 与 GRU 的组合逻辑

3.1 RGAT 在图上的消息传递

模型部分是这个项目的重头戏。溯源图是典型的异质图，节点类型有进程、文件、socket，边类型有“进程访问文件”“进程创建子进程”等。普通 GCN 或 GAT 只能处理同质图，而 RGAT（Relational Graph Attention Network）在注意力机制里加了一个“关系矩阵”，让每条边根据自己的类型影响邻居聚合的权重。

项目里model_RGAT.py实现了这个图注意力层，核心是下面这段逻辑：

import torch import torch.nn as nn import torch.nn.functional as F class RGATLayer(nn.Module): def __init__(self, in_dim, out_dim, num_relations, num_heads=4): super(RGATLayer, self).__init__() self.num_heads = num_heads self.num_relations = num_relations self.out_dim = out_dim # 每个关系都有自己的线性变换权重 self.relation_weights = nn.ModuleList([ nn.Linear(in_dim, out_dim * num_heads) for _ in range(num_relations) ]) self.attn = nn.Parameter(torch.zeros(num_heads, out_dim * 2)) def forward(self, x, edge_index, edge_type): # x: [num_nodes, in_dim] # edge_index: [2, num_edges] # edge_type: [num_edges] num_nodes = x.size(0) head_outputs = [] for head in range(self.num_heads): # 按每条边的关系类型选取对应的线性变换 transformed = torch.zeros(num_nodes, self.out_dim) for rel in range(self.num_relations): mask = edge_type == rel if mask.sum() > 0: subset_x = x[edge_index[0, mask]] transformed[edge_index[0, mask]] += self.relation_weights[rel](subset_x) # 注意力系数计算（简化版） attn_score = torch.matmul(transformed, self.attn[head]) # ... 这里做 softmax 和邻居聚合 head_outputs.append(attn_score) return torch.stack(head_outputs, dim=1).mean(dim=1)

这个简化版的 RGAT 层每个关系类型都有自己的Linear，相当于不同语义的边最后能投影到同一个向量空间里。实际项目中为了控制参数量，通常会共享一部分参数，但核心思想不变。需要注意的是edge_type必须和边列表一一对应，如果预处理阶段把边类型丢了，这里训练时会直接报维度不匹配的错。

3.2 GRU 对节点序列的时序建模

有了图注意力层之后，每个节点在每个时间窗口都有了一个向量表示。但 APT 攻击是时序过程，单纯把图卷积的输出拼在一起会丢失先后顺序。项目里用 GRU 来建模这种时间演变——把同一个节点在不同快照里的表示按时间顺序送进 GRU，输出最后一个时刻的隐藏状态作为该节点的最终表示。

model_RGAT_GRU.py里是这样组合的：

class RGAT_GRU(nn.Module): def __init__(self, in_dim, hidden_dim, num_relations, num_layers=2): super(RGAT_GRU, self).__init__() self.rgat = RGATLayer(in_dim, hidden_dim, num_relations) self.gru = nn.GRU(hidden_dim, hidden_dim, num_layers=num_layers, batch_first=True) def forward(self, snapshots, edge_indices, edge_types): # snapshots: 多个时间窗口的节点特征列表 node_reprs = [] for x, ei, et in zip(snapshots, edge_indices, edge_types): node_reprs.append(self.rgat(x, ei, et)) # 将每个快照的节点表示堆叠成序列 seq = torch.stack(node_reprs, dim=1) # [num_nodes, num_snapshots, hidden_dim] out, _ = self.gru(seq) return out[:, -1, :]

这里有个容易被忽略的细节：每个时间窗口内的节点集合可能不一样，比如某个进程只在某个窗口出现过。实际处理时通常会做一个全局节点对齐，把不存在的节点补成零向量，否则torch.stack会失败。项目里是在数据处理阶段先把所有时间窗口的节点并集求出来，再给每个窗口建一个“节点到行号”的映射表。

3.3 模型参数配置与训练命令

训练入口在main.py和main_GRU.py，区别在于是用纯 RGAT 还是 RGAT+GRU。核心超参数包括：

启动训练的命令很简单：

python main_GRU.py --dataset darpa --data_path ./data/darpa_tc5_cadets --epochs 200 --lr 0.001 --time_window 60

跑 StreamSpot 就把--dataset streamspot换掉。训练过程中会打印每个 epoch 的 loss 和验证集 F1，建议用--patience 20开早停，避免过拟合。我第一次没开早停，训练到第 180 轮的时候验证集分数反而往下掉，白等了半小时。

4. 跑通项目：环境准备、训练与评估的完整步骤

4.1 环境依赖与目录结构

这份源码是纯 Python 写的，依赖主要集中在 PyTorch 和它的图学习库。建议用 Conda 建一个干净环境：

conda create -n apt_detect python=3.8 conda activate apt_detect pip install torch==1.13.1 torch-geometric==2.2.0 pandas numpy scikit-learn

如果不想用 PyTorch Geometric，项目里也提供了纯 PyTorch 实现的 RGAT 层，那就不需要额外装torch_geometric。我建议先按裸 PyTorch 跑通，再考虑加几何库，这样排错容易。

目录结构保持源码原始布局就行：

. ├── main.py # 纯 RGAT 训练入口 ├── main_GRU.py # RGAT + GRU 训练入口 ├── model.py # 基础模型定义 ├── model_RGAT.py # RGAT 层定义 ├── model_RGAT_GRU.py # RGAT+GRU 组合模型 ├── streamspot.py # StreamSpot 数据解析 ├── streamspot_RGAT.py # StreamSpot 训练脚本 ├── darpa_cadets_RGAT.py # DARPA 训练脚本 ├── darpa_cadets_RGAT_GRU.py # DARPA 训练脚本（带 GRU） ├── data/ # 数据集目录 └── readme.md # 项目说明

4.2 训练与测试命令

数据放好之后，先跑一个最简命令验证环境：

python main.py --dataset streamspot --data_path ./data/streamspot --quick_test

项目里如果有--quick_test这种参数就直接用，没有的话可以自己注释掉训练循环，只跑一个 batch 的前向传播，检查有没有维度报错。这个习惯能帮你省掉大量浪费时间。

正式训练 DARPA 数据集时，我建议先用少量日志文件试跑，因为 DARPA 原始数据一天就几个 GB，全量跑非常吃内存。项目里如果提供了合并数据的脚本，那就用；没提供的话自己写个抽样函数，只取前 10 万条事件构建图。

4.3 评估指标怎么看

训练结束后，模型会输出在测试集上的混淆矩阵和 F1 分数。APT 检测场景里不能光看 accuracy，因为正常样本远多于异常样本，准确率再高也可能是把所有样本都判成了正常。重点看 precision 和 recall 的平衡：

• precision 高：检测到的攻击里真正的攻击比例高，误报少。
• recall 高：真实攻击被找到的比例高，漏报少。

安全场景通常更在意 recall，因为漏掉一次 APT 攻击比误报十次严重得多。项目默认的评估指标是 F1，如果你想调成偏向 recall，可以在测试脚本里加一个recall权重，或者降低分类阈值。

5. 避坑与排查：我在复现时踩过的五个坑

5.1 数据解析字段对不上，模型训练不收敛

现象：训练了 50 个 epoch，loss 一直稳定在某个值附近，精确率和召回率全是 0。

原因：数据集解析脚本里字段名写错了，导致大量边被过滤掉，图退化成孤立的点集合，模型学不到任何结构信息。

解决：先写一个数据检查函数，打印前 10 条边的src、dst、op，人工确认是不是符合常识。比如 DARPA 的op字段可能是open,read,write，如果解析出来全是unknown，大概率字段映射错了。

5.2 内存溢出，Process finished with exit code 137

现象：加载 DARPA 全量数据时，Python 进程直接被操作系统杀掉。

原因：DARPA 的日志文件是逐行 JSON，用readlines()一次性读入会导致内存瞬间爆炸。

解决：改成流式读取，逐行处理，不要把所有事件都存进 list。我一般先用wc -l看文件行数，超过 50 万行就一定会做数据抽样。也可以把图构建过程写成生成器，每处理 1 万条事件就释放一次内存。

5.3 节点 ID 映射错乱，图结构对不上

现象：训练时edge_index里的编号超出了x的行数，直接报索引越界。

原因：不同时间窗口的节点编号不统一，有的窗口新出现了进程，没有和全局编号对齐。

解决：在数据预处理阶段统一构建全局节点字典，先扫描所有窗口的原始节点名，再分配整数 ID。不要在每个窗口里单独用LabelEncoder，那个东西只能保证局部一致。

5.4 GRU 输入序列长度不一致，stack 报错

现象：torch.stack报Expected tensor for argument 1... to have the same size。

原因：不同时间窗口的节点数量不一样，每个快照的特征矩阵行数不同，而 GRU 要求所有序列等长。

解决：取所有窗口节点的并集作为统一节点集合，缺少的节点特征补零向量。代码里需要在构图时先用一个set收集所有节点，再为每个窗口生成固定大小的特征矩阵。

5.5 训练结果忽高忽低，同一次代码跑两遍分数不一样

现象：相同参数下两次训练 F1 差 10 个百分点以上。

原因：设置了随机种子，但 PyTorch 的某些算子在不固定种子时会有随机初始化，导致模型初始化参数不同。

解决：在所有入口文件开头固定torch.manual_seed(0)和np.random.seed(0)，还不行就把cudnn.deterministic = True也加上。这个坑最容易被人忽略，尤其是对比实验的时候，不固定种子的话你得不出可靠结论。

6. 进阶：换一个数据集、调一个阈值，验证你的模型没有过拟合

项目跑通之后，我建议你做两件额外的事：一是把 DARPA 和 StreamSpot 两个数据集交叉验证，二是画一条精确率-召回率曲线来挑阈值。

交叉验证很简单：用streamspot_RGAT.py训练出来的模型权重，直接加载到 DARPA 的测试脚本里做前向推理。因为两个数据集的图结构差异很大，如果模型在 StreamSpot 上表现好、在 DARPA 上立刻崩掉，说明它严重过拟合了某个数据集的特定模式，这个结论对答辩很有价值。

调阈值这件事，我用一段小代码记录一下：

from sklearn.metrics import precision_recall_curve # y_scores 是模型对测试样本输出的异常概率 precision, recall, thresholds = precision_recall_curve(y_true, y_scores) # 找一个 recall >= 0.9 且 precision 最高的阈值 best_threshold = None best_precision = 0 for thr, prec, rec in zip(thresholds, precision[:-1], recall[:-1]): if rec >= 0.9 and prec > best_precision: best_precision = prec best_threshold = thr print(f"Best threshold: {best_threshold:.4f}, precision: {best_precision:.4f}, recall >= 0.9")

这段代码的原理很简单：模型输出的概率在 0 和 1 之间，默认以 0.5 为界判断是否攻击。但实际场景里你宁可多报几个误报，也不希望漏掉真实攻击，所以把阈值降下来，用更高的 recall 来保护安全底线。项目测试文件里如果写死了 0.5，你可以把测试脚本的判别部分改成从这个阈值变量读取。

从那以后，我每次跑完模型都会先画一遍 PR 曲线，看看默认阈值是不是真的合适，再决定要不要调整。这个习惯帮我省掉了至少三次答辩时被老师追问“你的阈值怎么定的”的尴尬。希望这份源码也能帮你把 APT 检测这条链路完整跑通，少踩几个我踩过的坑。

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