The repository of ET-BERT, a network traffic classification model on encrypted traffic.
ET-BERT is a method for learning datagram contextual relationships from encrypted traffic, which could be directly applied to different encrypted traffic scenarios and accurately identify classes of traffic. First, ET-BERT employs multi-layer attention in large scale unlabelled traffic to learn both inter-datagram contextual and inter-traffic transport relationships. Second, ET-BERT could be applied to a specific scenario to identify traffic types by fine-tuning the labeled encrypted traffic on a small scale.
The work is introduced in the 31st The Web Conference:
Xinjie Lin, Gang Xiong, Gaopeng Gou, Zhen Li, Junzheng Shi and Jing Yu. 2022. ET-BERT: A Contextualized Datagram Representation with Pre-training Transformers for Encrypted Traffic Classification. In Proceedings of The Web Conference (WWW) 2022, Lyon, France. Association for Computing Machinery.
Note: this code is based on UER-py. Many thanks to the authors.
- Python >= 3.6
- torch >= 1.1
- six >= 1.12.0
- scapy == 2.4.4
- numpy == 1.19.2
- shutil, random, json, pickle, binascii, flowcontainer
- argparse
- packaging
- wireshark
- SplitCap
- scikit-learn
- For the mixed precision training you will need apex from NVIDIA
- For the pre-trained model conversion (related with TensorFlow) you will need TensorFlow
- For the tokenization with wordpiece model you will need WordPiece
- For the use of CRF in sequence labeling downstream task you will need pytorch-crf
The real-world TLS 1.3 dataset is collected from March to July 2021 on China Science and Technology Network (CSTNET). For privacy considerations, we only release the anonymous data (see in CSTNET-TLS 1.3).
Other datasets we used for comparison experiments are publicly available, see the paper for more details. If you want to use your own data, please check if the data format is the same as datasets/cstnet-tls1.3/
and specify the data path in data_preprocess/
.
You can now use ET-BERT directly through the pre-trained model or download via:
wget -O pretrained_model.bin https://drive.google.com/file/d/1r1yE34dU2W8zSqx1FkB8gCWri4DQWVtE/view?usp=sharing
After obtaining the pre-trained model, ET-BERT could be applied to the spetic task by fine-tuning with labeled network traffic:
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
The default path of the fine-tuned classifier model is models/finetuned_model.bin
. Then you can do inference with the fine-tuned model:
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
To reproduce the steps necessary to pre-train ET-BERT on network traffic data, follow the following steps:
- Run
data_preprocess/main.py
to generate the encrypted traffic corpus or directly use the generated corpus incorpora/
. Note you'll need to change the file paths and some configures at the top of the file. - Run
preprocess.py
to pre-process the encrypted traffic burst corpus.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
- Run
datasets/main.py
to generate the data for downstream tasks.
To reproduce the steps necessary to finetune ET-BERT on labeled data, run pretrain.py
to pre-train.
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
To see an example of how to use ET-BERT for the encrypted traffic classification tasks, go to the Using ET-BERT and run_classifier.py
script in the fine-tuning
folder.
Note: you'll need to change the path in programes.
If you are using the work (e.g. pre-trained model) in ET-BERT for academic work, please cite the paper published in WWW 2022:
@inproceedings{Lin2022etbert,
author = {Lin, Xinjie and Xiong, Gang and Gou, Gaopeng and Li, Zhen and Shi, Junzheng and Yu, Jing},
title = {ET-BERT: A Contextualized Datagram Representation with Pre-training Transformers for Encrypted Traffic Classification},
year = {2022},
month = apr,
isbn = {9781450390965},
publisher = {Association for Computing Machinery},
booktitle = {Proceedings of The Web Conference 2022},
numpages = {10},
}
Please post a Github issue if you have any questions.