ET-BERT：一种基于预训练Transformer的加密流量分类上下文数据报表示

注：

• ⭐ 如果您喜欢该项目，请留下一个星标！ ⭐
• 如果您发现任何错误 / 不当 / 过时的内容，请考虑提交一个问题或拉取请求。

ET-BERT代码库，一种针对加密流量的网络流量分类模型。

ET-BERT是一种从加密流量中学习数据报上下文关系的方法，可以直接应用于不同的加密流量场景，并准确识别流量类别。首先，ET-BERT在大规模未标记流量上采用多层注意力机制，学习数据报间以及不同流量传输之间的上下文关系。其次，通过在小规模标注加密流量上进行微调，ET-BERT可以应用于特定场景以识别流量类型。

该工作已在*第31届万维网大会*上发表：

Xinjie Lin, Gang Xiong, Gaopeng Gou, Zhen Li, Junzheng Shi 和 Jing Yu. 2022. ET-BERT：一种基于预训练Transformer的加密流量分类上下文数据报表示。载于《万维网大会》（WWW）2022年会议论文集，法国里昂。计算机协会。

注：此代码基于UER-py。非常感谢原作者。

目录

• 要求
• 数据集
• 使用ET-BERT
• 复现ET-BERT
• 引用
• 联系
• 赞助

要求

• Python >= 3.6
• CUDA: 11.4
• GPU: Tesla V100S
• torch >= 1.1
• six >= 1.12.0
• scapy == 2.4.4
• numpy == 1.19.2
• shutil, random, json, pickle, binascii, flowcontainer
• argparse
• packaging
• tshark
• SplitCap
• scikit-learn
• 对于混合精度训练，您需要NVIDIA的apex工具包。
• 对于预训练模型的转换（与TensorFlow相关），您需要TensorFlow。
• 对于使用wordpiece模型进行分词，您需要WordPiece。
• 对于在序列标注下游任务中使用CRF，您需要pytorch-crf。

数据集

真实世界的TLS 1.3数据集收集于2021年3月至7月期间，来自中国科学技术网（CSTNET）。出于隐私考虑，我们仅发布匿名数据（详见CSTNET-TLS 1.3）。

我们用于对比实验的其他数据集均为公开可用，更多细节请参阅论文。如果您希望使用自己的数据，请确保数据格式与datasets/cstnet-tls1.3/一致，并在data_process/中指定数据路径。

使用ET-BERT

您现在可以直接使用预训练的模型，或者通过以下命令下载：

wget -O pretrained_model.bin https://drive.google.com/file/d/1r1yE34dU2W8zSqx1FkB8gCWri4DQWVtE/view?usp=sharing

获得预训练模型后，可以通过在数据包级别对标注的网络流量进行微调，将ET-BERT应用于特定任务：

python3 fine-tuning/run_classifier.py --pretrained_model_path models/pre-trained_model.bin \
                                   --vocab_path models/encryptd_vocab.txt \
                                   --train_path datasets/cstnet-tls1.3/packet/train_dataset.tsv \
                                   --dev_path datasets/cstnet-tls1.3/packet/valid_dataset.tsv \
                                   --test_path datasets/cstnet-tls1.3/packet/test_dataset.tsv \
                                   --epochs_num 10 --batch_size 32 --embedding word_pos_seg \
                                   --encoder transformer --mask fully_visible \
                                   --seq_length 128 --learning_rate 2e-5

微调后的分类器模型默认保存路径为models/finetuned_model.bin。随后您可以使用微调后的模型进行推理：

python3 inference/run_classifier_infer.py --load_model_path models/finetuned_model.bin \
                                          --vocab_path models/encryptd_vocab.txt \
                                          --test_path datasets/cstnet-tls1.3/packet/nolabel_test_dataset.tsv \
                                          --prediction_path datasets/cstnet-tls1.3/packet/prediction.tsv \
                                          --labels_num 120 \
                                          --embedding word_pos_seg --encoder transformer --mask fully_visible

复现ET-BERT

预处理

要复现ET-BERT在网络流量数据上进行预训练所需的步骤，请按照以下步骤操作：

1. 运行vocab_process/main.py生成加密流量语料库，或直接使用corpora/中已生成的语料库。请注意，您需要更改文件顶部的文件路径和部分配置。
2. 运行main/preprocess.py对加密流量突发语料库进行预处理。

      python3 preprocess.py --corpus_path corpora/encrypted_traffic_burst.txt \
                            --vocab_path models/encryptd_vocab.txt \
                            --dataset_path dataset.pt --processes_num 8 --target bert
   

3. 运行data_process/main.py生成下游任务所需的数据，如果存在需要处理的pcap格式数据集。此过程包括两个步骤。首先，在datasets/main.py:54中设置splitcap=True来拆分pcap文件，并将其保存为npy格式的数据集。然后是生成微调数据的过程。如果您使用共享数据集，则需要在dataset_save_path下创建名为dataset的文件夹，并将数据集复制到此处。

预训练

要复现在标注数据上微调 ET-BERT 所需的步骤，请运行 pretrain.py 进行预训练。如果希望在已预训练的模型基础上继续训练，可以增加参数 --pretrained_model_path。

   python3 pre-training/pretrain.py --dataset_path dataset.pt --vocab_path models/encryptd_vocab.txt \
                       --output_model_path models/pre-trained_model.bin \
                       --world_size 8 --gpu_ranks 0 1 2 3 4 5 6 7 \
                       --total_steps 500000 --save_checkpoint_steps 10000 --batch_size 32 \
                       --embedding word_pos_seg --encoder transformer --mask fully_visible --target bert

下游任务上的微调

要查看如何将 ET-BERT 用于加密流量分类任务的示例，请参阅 使用 ET-BERT 以及 fine-tuning 文件夹中的 run_classifier.py 脚本。

注意：您需要更改程序中的路径。

引用

如果您在学术工作中使用了 ET-BERT 中的工作（例如预训练模型），请引用发表于 WWW 2022 的论文：

@inproceedings{lin2022etbert,
  author    = {Xinjie Lin and
               Gang Xiong and
               Gaopeng Gou and
               Zhen Li and
               Junzheng Shi and
               Jing Yu},
  title     = {{ET-BERT:} {A} Contextualized Datagram Representation with Pre-training
               Transformers for Encrypted Traffic Classification},
  booktitle = {{WWW} '22: The {ACM} Web Conference 2022, Virtual Event, Lyon, France,
               April 25 - 29, 2022},
  pages     = {633--642},
  publisher = {{ACM}},
  year      = {2022}
}

联系方式

如果您有任何问题，请在 GitHub 上提交 issue。欢迎讨论新的想法、技术及改进！

请我喝杯咖啡

如果您觉得本项目有用并想表达感谢，欢迎请我喝杯咖啡！这将有助于保持我的动力和热情！

ET-BERT 快速上手指南

ET-BERT 是一个基于预训练 Transformer 的加密流量分类模型，能够学习数据报的上下文关系，直接应用于不同的加密流量场景并准确识别流量类型。

环境准备

系统要求

• 操作系统: Linux (推荐)
• Python: >= 3.6
• CUDA: 11.4
• GPU: 推荐 Tesla V100S 或同等算力显卡

前置依赖

请确保安装以下核心依赖库：

pip install torch>=1.1 six>=1.12.0 scapy==2.4.4 numpy==1.19.2 argparse packaging scikit-learn

其他可选依赖（根据需求安装）：

• 混合精度训练: 需安装 NVIDIA apex
• TensorFlow 模型转换: 需安装 tensorflow
• WordPiece 分词: 需安装 tokenizers (pip install tokenizers)
• CRF 序列标注: 需安装 pytorch-crf (pip install pytorch-crf)
• 工具链: 需安装 tshark 和 SplitCap 用于数据包处理。

安装步骤

1. 克隆项目代码

   git clone https://github.com/linwhitehat/et-bert.git
   cd et-bert
   

2. 下载预训练模型 您可以直接从 Google Drive 下载官方提供的预训练模型文件：

   wget -O pretrained_model.bin "https://drive.google.com/file/d/1r1yE34dU2W8zSqx1FkB8gCWri4DQWVtE/view?usp=sharing"
   

   注：如果下载速度慢，建议手动在浏览器打开链接下载后，将文件重命名为 pretrained_model.bin 并放入 models/ 目录。

3. 准备数据集 项目默认使用 CSTNET-TLS 1.3 数据集。如需使用自有数据，请确保数据格式与 datasets/cstnet-tls1.3/ 目录下的结构一致，并在 data_process/ 中配置路径。

基本使用

以下是使用预训练模型进行**微调（Fine-tuning）和推理（Inference）**的最简流程。

1. 微调模型

使用标注好的加密流量数据对预训练模型进行微调，以适应特定分类任务：

python3 fine-tuning/run_classifier.py --pretrained_model_path models/pre-trained_model.bin \
                                   --vocab_path models/encryptd_vocab.txt \
                                   --train_path datasets/cstnet-tls1.3/packet/train_dataset.tsv \
                                   --dev_path datasets/cstnet-tls1.3/packet/valid_dataset.tsv \
                                   --test_path datasets/cstnet-tls1.3/packet/test_dataset.tsv \
                                   --epochs_num 10 --batch_size 32 --embedding word_pos_seg \
                                   --encoder transformer --mask fully_visible \
                                   --seq_length 128 --learning_rate 2e-5

执行完成后，微调后的模型将保存为 models/finetuned_model.bin。

2. 模型推理

使用微调后的模型对未标注的流量数据进行分类预测：

python3 inference/run_classifier_infer.py --load_model_path models/finetuned_model.bin \
                                          --vocab_path models/encryptd_vocab.txt \
                                          --test_path datasets/cstnet-tls1.3/packet/nolabel_test_dataset.tsv \
                                          --prediction_path datasets/cstnet-tls1.3/packet/prediction.tsv \
                                          --labels_num 120 \
                                          --embedding word_pos_seg --encoder transformer --mask fully_visible

提示：运行前请根据实际文件路径修改上述命令中的 --*_path 参数。