ViktorAxelsen/MH-Net

# MH-Net AAAIC'25 主赛道论文《Revolutionizing Encrypted Traffic Classification with MH-Net: A Multi-View Heterogeneous Graph Model》的官方实现

 ## 环境配置 ``` # python==3.8 pip install torch==1.10.1+cu113 torchvision==0.11.2+cu113 torchaudio==0.10.1 -f https://download.pytorch.org/whl/cu113/torch_stable.html pip install dgl==1.0.0+cu113 -f https://data.dgl.ai/wheels/cu113/repo.html pip install scikit-learn pip install scapy ``` ## 预处理 MH-Net 的预处理与 [TFE-GNN](https://github.com/ViktorAxelsen/TFE-GNN) 类似。您可以选择跳过预处理部分。 ### 将 .pcap 转换为 .npz 文件 为了便于后续处理，我们从 .pcap 文件中提取信息并保存为 .npz 文件。 请参考 **config.py**，并在 **DIR_PATH_DICT** 中自定义您的 .pcap 路径。然后运行以下命令开始转换。 ``` # ISCX-VPN python pcap2npy.py --dataset iscx-vpn # ISCX-NonVPN python pcap2npy.py --dataset iscx-nonvpn # ISCX-TOR python pcap2npy.py --dataset iscx-tor # ISCX-NonTOR python pcap2npy.py --dataset iscx-nontor # CIC-Iot python pcap2npy.py --dataset ciciot ``` ### 构建多视图异构流量图 在运行以下命令之前，您可以参考 **config.py** 并自定义所有文件路径。然后运行以下命令开始构建。此外，我们仅提供 ciciot 数据集的完整路径作为参考，您可以自行修改其他路径。 您可以通过指定 "transform_length" 来生成不同比特长度的异构流量图。 $Note$：我们以 8 位作为基线，并将 8 位流量图转换为其他比特长度。 ``` # ISCX-VPN python preprocess.py --dataset iscx-vpn --transform_length 4 # ISCX-NonVPN python preprocess.py --dataset iscx-nonvpn --transform_length 4 # ISCX-TOR python preprocess.py --dataset iscx-tor --transform_length 4 # ISCX-NonTOR python preprocess.py --dataset iscx-nontor --transform_length 4 # CIC-IoT python preprocess.py --dataset ciciot --transform_length 4 ``` ## 训练 运行以下命令开始训练。 ``` # ISCX-VPN CUDA_VISIBLE_DEVICES="0" python train_new.py --dataset iscx-vpn --prefix exp_train --coe 0.5 --coe_graph 1.0 --seq_aug_ratio 0.6 --drop_edge_ratio 0.05 --drop_node_ratio 0.1 --K 15 --hp_ratio 0.5 --tau 0.07 --gtau 0.07 # ISCX-NonVPN CUDA_VISIBLE_DEVICES="0" python train_new.py --dataset iscx-nonvpn --prefix exp_train --coe 0.8 --coe_graph 0.4 --seq_aug_ratio 0.6 --drop_edge_ratio 0.05 --drop_node_ratio 0.1 --K 15 --hp_ratio 0.5 --tau 0.07 --gtau 0.07 # ISCX-TOR CUDA_VISIBLE_DEVICES="0" python train_new.py --dataset iscx-tor --prefix exp_train --coe 1.0 --coe_graph 0.4 --seq_aug_ratio 0.6 --drop_edge_ratio 0.05 --drop_node_ratio 0.1 --K 15 --hp_ratio 0.5 --tau 0.07 --gtau 0.07 # ISCX-NonTOR CUDA_VISIBLE_DEVICES="0" python train_new.py --dataset iscx-nontor --prefix exp_train --coe 1.0 --coe_graph 0.6 --seq_aug_ratio 0.6 --drop_edge_ratio 0.05 --drop_node_ratio 0.1 --K 15 --hp_ratio 0.5 --tau 0.07 --gtau 0.07 #CIC-IoT CUDA_VISIBLE_DEVICES="0" python train_new.py --dataset ciciot --prefix exp_train --coe 1.0 --coe_graph 0.6 --seq_aug_ratio 0.6 --drop_edge_ratio 0.05 --drop_node_ratio 0.1 --K 15 --hp_ratio 0.5 --tau 0.07 --gtau 0.07 ``` ## 评估 运行以下命令开始评估。 ``` # ISCX-VPN CUDA_VISIBLE_DEVICES="0" python test_new.py --dataset iscx-vpn --prefix exp_train # ISCX-NonVPN CUDA_VISIBLE_DEVICES="0" python test_new.py --dataset iscx-nonvpn --prefix exp_train # ISCX-TOR CUDA_VISIBLE_DEVICES="0" python test_new.py --dataset iscx-tor --prefix exp_train # ISCX-NonTOR CUDA_VISIBLE_DEVICES="0" python test_new.py --dataset iscx-nontor --prefix exp_train # CIC-IoT CUDA_VISIBLE_DEVICES="0" python test_new.py --dataset ciciot --prefix exp_train ``` ## 引用 ``` @article{MH-Net, title={Revolutionizing Encrypted Traffic Classification with MH-Net: A Multi-View Heterogeneous Graph Model}, author={Haozhen Zhang and Haodong Yue and Xi Xiao and Le Yu and Qing Li and Zhen Ling and Ye Zhang}, journal={arXiv preprint arXiv:2501.03279}, year={2025} } ```

标签：AAAI 2025, arXiv:2501.03279, DGL, ISCX-NonTOR, ISCX-NonVPN, ISCX-TOR, ISCX-VPN, MH-Net, .pcap转.npz, PyTorch, 凭据扫描, 加密流量分析, 加密流量分类, 图神经网络, 多视图学习, 多视图异构图模型, 异构图, 流量分类模型, 流量特征提取, 深度学习流量分类, 网络流量分类, 逆向工具