def validate_process(val_loader, model, criterion, device, print_freq):
losses = AverageMeter()
top1 = AverageMeter()
model.eval() # switch to evaluate mode
for i, (pcap, statistic, target) in enumerate(val_loader):
pcap = pcap.to(device)
statistic = statistic.to(device)
target = target.to(device)
with torch.no_grad():
output = model(pcap, statistic) # compute output
loss = criterion(output, target) # 计算验证集的 loss
# measure accuracy and record loss
prec1 = accuracy(output.data, target)
losses.update(loss.item(), pcap.size(0))
top1.update(prec1[0].item(), pcap.size(0))
if (i + 1) % print_freq == 0:
logger.info(
'Test: [{0}/{1}], Loss {loss.val:.4f} ({loss.avg:.4f}), Prec@1 {top1.val:.3f} ({top1.avg:.3f})'
.format(i, len(val_loader), loss=losses, top1=top1))
logger.info(' * Prec@1 {top1.avg:.3f}'.format(top1=top1))
return top1.avg
def validate_process(val_loader, model, device, print_freq):
top1 = AverageMeter()
model.eval() # switch to evaluate mode
for i, (pcap, statistic, target) in enumerate(val_loader):
pcap = (pcap / 255).to(device) # 也要归一化
statistic = statistic.to(device)
statistic = (statistic -
mean_val) / std_val # 首先需要对 statistic 的数据进行归一化
target = target.to(device)
with torch.no_grad():
output, _ = model(pcap, statistic) # compute output
# measure accuracy and record loss
prec1 = accuracy(output.data, target)
top1.update(prec1[0].item(), pcap.size(0))
if (i + 1) % print_freq == 0:
logger.info(
'Test: [{0}/{1}], Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.
format(i, len(val_loader), top1=top1))
logger.info(' * Prec@1 {top1.avg:.3f}'.format(top1=top1))
return top1.avg
def train_process(train_loader, model, alpha, criterion_c, criterion_r,
optimizer, epoch, device, print_freq):
"""训练一个 epoch 的流程
Args:
train_loader (dataloader): [description]
model ([type]): [description]
criterion_c ([type]): 计算分类误差
criterion_l ([type]): 计算重构误差
optimizer ([type]): [description]
epoch (int): 当前所在的 epoch
device (torch.device): 是否使用 gpu
print_freq ([type]): [description]
"""
c_loss = AverageMeter()
r_loss = AverageMeter()
losses = AverageMeter() # 在一个 train loader 中的 loss 变化
top1 = AverageMeter() # 记录在一个 train loader 中的 accuracy 变化
model.train() # 切换为训练模型
for i, (pcap, statistic, target) in enumerate(train_loader):
pcap = (pcap / 255).to(device) # 也要归一化
statistic = statistic.to(device)
statistic = (statistic -
mean_val) / std_val # 首先需要对 statistic 的数据进行归一化
target = target.to(device)
classific_result, fake_statistic = model(pcap, statistic) # 分类结果和重构结果
loss_c = criterion_c(classific_result, target) # 计算 分类的 loss
loss_r = criterion_r(statistic, fake_statistic) # 计算 重构 loss
loss = alpha * loss_c + loss_r # 将两个误差组合在一起
# 计算准确率, 记录 loss 和 accuracy
prec1 = accuracy(classific_result.data, target)
c_loss.update(loss_c.item(), pcap.size(0))
r_loss.update(loss_r.item(), pcap.size(0))
losses.update(loss.item(), pcap.size(0))
top1.update(prec1[0].item(), pcap.size(0))
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i + 1) % print_freq == 0:
logger.info(
'Epoch: [{0}][{1}/{2}], Loss {loss.val:.4f} ({loss.avg:.4f}), Loss_c {loss_c.val:.4f} ({loss_c.avg:.4f}), Loss_r {loss_r.val:.4f} ({loss_r.avg:.4f}), Prec@1 {top1.val:.3f} ({top1.avg:.3f})'
.format(epoch,
i,
len(train_loader),
loss=losses,
loss_c=c_loss,
loss_r=r_loss,
top1=top1))
def statisticFeature2JSON(folder_path):
"""将 folder_path 中所有的 session 计算统计特征, 并保存为 json 文件;
对于一些较大的 pcap 文件, 为了速度, 我们都只处理前 500000 个 packets
Args:
folder_path (str): 所在的路径
"""
pcap_statisticFeature = {}
featuresCalc = FeaturesCalc(min_window_size=1) # 初始化计算统计特征的类
for (root, _, files) in os.walk(folder_path):
logger.info('正在提取 {} 下 pcap 文件的统计特征'.format(root))
for Ufile in tqdm.tqdm(files):
pcapPath = os.path.join(root, Ufile) # 需要转换的pcap文件的完整路径
packets = rdpcap(pcapPath) # 读入 pcap 文件
if len(packets) < 500000: # 太大的 pcap 文件
features = featuresCalc.compute_features(
packets_list=packets) # 计算特征
else:
logger.info('正在处理 {} 文件'.format(pcapPath))
features = featuresCalc.compute_features(
packets_list=packets[:500000])
pcap_statisticFeature[Ufile] = features
# 将统计特征写入 json 文件
with open("statistic_features.json", "w") as f:
json.dump(pcap_statisticFeature, f)
logger.info('统计特征计算完成, 保存为 json 文件!')
logger.info('==========\n')
def data_loader(pcap_file,
statistic_file,
label_file,
trimed_file_len,
batch_size=256,
workers=1,
pin_memory=True):
"""读取处理好的 npy 文件, 并返回 pytorch 训练使用的 dataloader 数据
Args:
pcap_file (str): pcap 文件转换得到的 npy 文件的路径
statistic_file (str): 统计特征对应的 npy 文件路径
label_file (str): 上面的 pcap 文件对应的 label 文件的 npy 文件的路径
trimed_file_len (int): pcap 被裁剪成的长度
batch_size (int, optional): 默认一个 batch 有多少数据. Defaults to 256.
workers (int, optional): 处理数据的进程的数量. Defaults to 1.
pin_memory (bool, optional): 锁页内存, 如果内存较多, 可以设置为 True, 可以加快 GPU 的使用. Defaults to True.
Returns:
DataLoader: pytorch 训练所需要的数据
"""
# 载入 npy 数据
pcap_data = np.load(pcap_file) # 获得 pcap 文件
statistic_data = np.load(statistic_file) # 获得 统计特征
label_data = np.load(label_file) # 获得 label 数据
# 将 npy 数据转换为 tensor 数据
pcap_data = torch.from_numpy(pcap_data.reshape(-1, 1,
trimed_file_len)).float()
statistic_data = torch.from_numpy(statistic_data).float()
label_data = torch.from_numpy(label_data).long()
logger.info('pcap 文件大小, {}; statistic 文件大小, {}; label 文件大小: {}'.format(
pcap_data.shape, statistic_data.shape, label_data.shape))
# 将 tensor 数据转换为 Dataset->Dataloader
res_dataset = torch.utils.data.TensorDataset(pcap_data, statistic_data,
label_data) # 合并数据
res_dataloader = torch.utils.data.DataLoader(
dataset=res_dataset,
batch_size=batch_size,
shuffle=True,
pin_memory=pin_memory,
num_workers=1 # set multi-work num read data
)
return res_dataloader
def save_pcap2npy(pcap_dict,
file_name,
statistic_feature_json,
label2index={}):
"""将 pcap 文件分为训练集和测试集, 并进行保存.
=> pcap 数据从 二进制 转换为 十进制, 调用 getIntfrom_pcap 函数
=> 从 json 文件读取每一个 pcap 的统计特征
=> 将数据保存在 data 中, data = [[pcap, statistic, label], [], ..]
=> 将 data 数据保存为 npy 文件
Args:
pcap_dict (dict): 见函数 get_train_test 的输出
file_name (str): 最后保存的文件的名称
"""
with open(statistic_feature_json, "r") as json_read:
statistic_feature_dict = json.load(json_read) # 获取每个 pcap 的统计特征
data = []
index = 0
for label, pcap_file_list in pcap_dict.items():
logger.info('模式, {}, 正在将 {} 的 pcap 保存为 npy'.format(file_name, label))
if label not in label2index:
label2index[label] = index
index = index + 1
for pcap_file in pcap_file_list:
pcap_file_name = os.path.normpath(pcap_file).split('\\')[-1]
pcap_content = getIntfrom_pcap(
pcap_file) # 返回 pcap 的 十进制 内容, 这里 pcap_file 是 pcap 文件的路径名
statistic_feature = statistic_feature_dict[
pcap_file_name] # 得到统计特征
data.append([pcap_content, statistic_feature,
label2index[label]]) # 将 pcap 和 label 添加到 data 中去
np.random.shuffle(data) # 对数据进行打乱
pcap_data = np.array([i[0] for i in data]) # raw pcap 减裁
statistic_data = np.array([i[1] for i in data]) # statistic data
y = np.array([i[2] for i in data]) # label
logger.info('数据的大小, {}; 统计特征的大小, {}; 标签的大小, {};'.format(
pcap_data.shape, statistic_data.shape, y.shape)) # 打印数据大小
np.save('{}-pcap.npy'.format(file_name), pcap_data)
np.save('{}-statistic.npy'.format(file_name), statistic_data)
np.save('{}-labels.npy'.format(file_name), y)
logger.info('将 {} 的数据保存为 npy !'.format(file_name))
logger.info('==========\n')
return label2index
def train_process(train_loader, model, criterion, optimizer, epoch, device, print_freq):
"""训练一个 epoch 的流程
Args:
train_loader (dataloader): [description]
model ([type]): [description]
criterion ([type]): [description]
optimizer ([type]): [description]
epoch (int): 当前所在的 epoch
device (torch.device): 是否使用 gpu
print_freq ([type]): [description]
"""
losses = AverageMeter() # 在一个 train loader 中的 loss 变化
top1 = AverageMeter() # 记录在一个 train loader 中的 accuracy 变化
model.train() # 切换为训练模型
for i, (pcap, statistic, target) in enumerate(train_loader):
pcap = pcap.to(device)
statistic = statistic.to(device)
target = target.to(device)
output = model(pcap, statistic) # 得到模型预测结果
loss = criterion(output, target) # 计算 loss
# 计算准确率, 记录 loss 和 accuracy
prec1 = accuracy(output.data, target)
losses.update(loss.item(), pcap.size(0))
top1.update(prec1[0].item(), pcap.size(0))
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i+1) % print_freq == 0:
logger.info('Epoch: [{0}][{1}/{2}], Loss {loss.val:.4f} ({loss.avg:.4f}), Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(
epoch, i, len(train_loader), loss=losses, top1=top1))
def preprocess_pipeline():
"""对流量进行预处理, 处理流程为:
0. 将所有流量文件转移到新的文件夹, 这时候没有细分, 就是所有文件进行转移
1. 接着将 pcapng 文件转换为 pcap 文件
2. 接着将不同的流量新建文件夹, 分成不同的类别, pcap transfer
3. 接着将 pcap 文件按照五元组分为不同的 session, 使用 SplitCap 来完成 (这一步可以选择, 提取 all 或是 L7)
计算
4. 对 session 进行处理, 匿名化处理, ip, mac, port
5. 将所有的 pcap 转换为一样的大小, 转换前统计一下原始 session 的大小
6. 对于每一类的文件, 划分训练集和测试集, 获得每一类的所有的 pcap 的路径
7. 将所有的文件, 最终保存为 npy 的格式
"""
cfg = setup_config() # 获取 config 文件
logger.info(cfg)
transfer_pcap(cfg.pcap_path.raw_pcap_path,
cfg.pcap_path.new_pcap_path) # 转移文件
pcapng_to_pcap(cfg.pcap_path.new_pcap_path) # 将 pcapng 转换为 pcap
pcap_transfer(
cfg.pcap_path.new_pcap_path,
cfg.pcap_path.new_pcap_path) # 将文件放在指定文件夹中, 这里新的文件夹查看 yaml 配置文件
pcap_to_session(cfg.pcap_path.new_pcap_path,
cfg.tool_path.splitcap_path) # 将 pcap 转换为 session
statisticFeature2JSON(
cfg.pcap_path.new_pcap_path) # 计算 pcap 的统计特征 (特征可以只计算一次, 后面就不需要再运行了)
anonymize(cfg.pcap_path.new_pcap_path) # 对指定文件夹内的所有 pcap 进行匿名化处理
pcap_trim(
cfg.pcap_path.new_pcap_path,
cfg.train.TRIMED_FILE_LEN) # 将所有的 pcap 转换为一样的大小, 同时统计原始 session 的长度
train_dict, test_dict = get_train_test(
cfg.pcap_path.new_pcap_path,
cfg.train.train_size) # 返回 train 和 test 的 dict
label2index = save_pcap2npy(
train_dict, 'train',
cfg.pcap_path.statistic_feature) # 保存 train 的 npy 文件
save_pcap2npy(test_dict, 'test', cfg.pcap_path.statistic_feature,
label2index) # 保存 test 的 npy 文件
logger.info('index 与 label 的关系, {}'.format(label2index))
def get_tensor_data(pcap_file, statistic_file, label_file, trimed_file_len):
"""读取处理好的 npy 文件, 并返回 pytorch 训练使用的 dataloader 数据
Args:
pcap_file (str): pcap 文件转换得到的 npy 文件的路径
statistic_file (str): 统计特征转换得到的 npy 文件的路径
label_file (str): 上面的 pcap 文件对应的 label 文件的 npy 文件的路径
trimed_file_len (int): pcap 被裁剪成的长度
"""
# 载入 npy 数据
pcap_data = np.load(pcap_file) # 获得 pcap 文件
statistic_data = np.load(statistic_file)
label_data = np.load(label_file) # 获得 label 数据
# 将 npy 数据转换为 tensor 数据
pcap_data = torch.from_numpy(pcap_data.reshape(-1, 1,
trimed_file_len)).float()
statistic_data = torch.from_numpy(statistic_data).float()
label_data = torch.from_numpy(label_data).long()
logger.info(
'导入 Tensor 数据, pcap 文件大小, {}; statistic 大小, {}; label 文件大小: {}'.format(
pcap_data.shape, statistic_data.shape, label_data.shape))
return pcap_data, statistic_data, label_data
def pcap_transfer(before_folder_path, new_pcap_path):
"""将原始的 pcap 文件从旧的文件夹, 转移到新的文件夹, 并进行好分类
Args:
before_folder_path (str): 旧的文件夹的名称
new_pcap_path (str): 新的文件夹的名称
"""
for pcap_type in PCAP_LABEL_DICT:
logger.info('开始移动 {} 类型的 pcap 文件'.format(pcap_type))
folder_path = os.path.join(new_pcap_path, pcap_type) # 新的文件夹
os.makedirs(folder_path, exist_ok=True) # 新建新的文件夹
for pcap_name in PCAP_LABEL_DICT[pcap_type]:
pcap_name_file = '{}.pcap'.format(pcap_name)
src_path = os.path.join(before_folder_path,
pcap_name_file) # pcap 文件的原始地址
dts_path = os.path.join(folder_path, pcap_name_file)
os.rename(src_path, dts_path) # 移动文件
logger.info('文件移动, {} --> {}'.format(src_path, dts_path))
logger.info('文件移动完毕.')
logger.info('============\n')
def transfer_pcap(before_path, after_path):
"""将 before_path 中的所以文件转移到 after_path 中去
Args:
before_path (str): 转移之前的文件路径
after_path (str): 转移之后的文件路径
"""
ignore_list = [
'youtubeHTML5_1.pcap', 'torFacebook.pcap', 'torGoogle.pcap',
'torTwitter.pcap'
] # 不需要转移的 pcap 文件
os.makedirs(after_path, exist_ok=True) # 新建目标目录
for file in os.listdir(before_path):
src_path = os.path.join(before_path, file)
dst_path = os.path.join(after_path, file)
if file in ignore_list:
logger.info('忽略文件 {}'.format(src_path))
else:
logger.info('开始转移文件, {} --> {}'.format(src_path, dst_path))
shutil.copy(src_path, dst_path)
logger.info('文件全部转移完成')
logger.info('=============\n')
def get_file_path(folder_path):
"""获得 folder_path 下 pcap 文件的路径, 以 dict 的形式返回.
返回的包含每个大类(Chat, Email), 下每个小类(AIMchat1, aim_chat_3a), 中 pcap 的文件路径.
返回数据类型如下所示:
{
'Chat': {
'AIMchat1': ['D:\\Traffic-Classification\\data\\preprocess_data\\Chat\\AIMchat1\\AIMchat1.pcap.TCP_131-202-240-87_13393_178-237-24-202_443.pcap', ...]
'aim_chat_3a': [...],
...
},
'Email': {
'email1a': [],
...
},
...
}
Args:
folder_path (str): 包含 pcap 文件的根目录名称
"""
pcap_dict = {}
for (root, _, files) in os.walk(folder_path):
if len(files) > 0:
logger.info('正在记录 {} 下的 pcap 文件'.format(root))
folder_name_list = os.path.normpath(root).split(
os.sep
) # 将 'D:\Traffic-Classification\data\preprocess_data' 返回为列表 ['D:', 'Traffic-Classification', 'data', 'preprocess_data']
top_category, second_category = folder_name_list[
-2], folder_name_list[-1]
if top_category not in pcap_dict:
pcap_dict[top_category] = {}
if second_category not in pcap_dict[top_category]:
pcap_dict[top_category][second_category] = []
for Ufile in files:
pcapPath = os.path.join(root, Ufile) # 需要转换的pcap文件的完整路径
pcap_dict[top_category][second_category].append(pcapPath)
logger.info('将所有的 pcap 文件整理为 dict !')
logger.info('==========\n')
return pcap_dict
def pcap_to_session(pcap_folder, splitcap_path):
"""将 pcap 文件转换为 session 文件, 将 pcap_folder 下所有的 pcap 文件进行转换
Args:
pcap_folder (str): 放置 pcap 文件的路径
splitcap_path (str): splitcap.exe 工具所在的路径
"""
splitcap_path = os.path.normpath(splitcap_path) # 处理成 windows 下的路径格式
for (root, _, files) in os.walk(pcap_folder):
# root 是根目录
# dirs 是在 root 目录下的文件夹, 返回的是一个 list;
# files 是在 root 目录下的文件, 返回的是一个 list, 所以 os.path.join(root, files) 返回的就是 files 的路径
for Ufile in files:
pcap_file_path = os.path.join(root, Ufile) # pcap 文件的完整路径
pcap_name = Ufile.split('.')[0] # pcap 文件的名字
pcap_suffix = Ufile.split('.')[1] # 文件的后缀名
try:
assert pcap_suffix == 'pcap'
except:
logger.warning('查看 pcap 文件的后缀')
assert pcap_suffix == 'pcap'
os.makedirs(os.path.join(root, pcap_name), exist_ok=True) # 新建文件夹
prog = subprocess.Popen(
[
splitcap_path, "-p", "100000", "-b", "100000", "-r",
pcap_file_path, "-o",
os.path.join(root, pcap_name)
], # 只提取应用层可以加上, "-y", "L7"
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
shell=True)
_, _ = prog.communicate()
# logger.info(err.decode('GB2312'))
os.remove(pcap_file_path) # 删除原始的 pcap 文件
logger.info('处理完成文件 {}'.format(Ufile))
logger.info('完成 pcap 转换为 session.')
logger.info('============\n')
def CENTIME_train_pipeline(alpha):
cfg = setup_config() # 获取 config 文件
logger.info(cfg)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger.info('是否使用 GPU 进行训练, {}'.format(device))
model_path = os.path.join(cfg.train.model_dir,
cfg.train.model_name) # 模型的路径
model = resnet_AE(model_path, pretrained=False,
num_classes=12).to(device) # 初始化模型
criterion_c = nn.CrossEntropyLoss() # 分类用的损失函数
criterion_r = nn.L1Loss() # 重构误差的损失函数
optimizer = optim.Adam(model.parameters(), lr=cfg.train.lr) # 定义优化器
logger.info('成功初始化模型.')
train_loader = data_loader(
pcap_file=cfg.train.train_pcap,
label_file=cfg.train.train_label,
statistic_file=cfg.train.train_statistic,
trimed_file_len=cfg.train.TRIMED_FILE_LEN) # 获得 train dataloader
test_loader = data_loader(
pcap_file=cfg.train.test_pcap,
label_file=cfg.train.test_label,
statistic_file=cfg.train.test_statistic,
trimed_file_len=cfg.train.TRIMED_FILE_LEN) # 获得 train dataloader
logger.info('成功加载数据集.')
best_prec1 = 0
for epoch in range(cfg.train.epochs):
adjust_learning_rate(optimizer, epoch, cfg.train.lr) # 动态调整学习率
train_process(train_loader, model, alpha, criterion_c, criterion_r,
optimizer, epoch, device, 80) # train for one epoch
prec1 = validate_process(test_loader, model, device,
20) # evaluate on validation set
# remember the best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
# 保存最优的模型
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict()
}, is_best, model_path)
# 下面进入测试模式, 计算每个类别详细的准确率
logger.info('进入测试模式.')
model = resnet_AE(model_path, pretrained=True,
num_classes=12).to(device) # 加载最好的模型
index2label = {j: i
for i, j in cfg.test.label2index.items()
} # index->label 对应关系
label_list = [index2label.get(i) for i in range(12)] # 12 个 label 的标签
pcap_data, statistic_data, label_data = get_tensor_data(
pcap_file=cfg.train.test_pcap,
statistic_file=cfg.train.test_statistic,
label_file=cfg.train.test_label,
trimed_file_len=cfg.train.TRIMED_FILE_LEN) # 将 numpy 转换为 tensor
pcap_data = (pcap_data / 255).to(device) # 流量数据
statistic_data = (statistic_data.to(device) -
mean_val) / std_val # 对数据做一下归一化
y_pred, _ = model(pcap_data, statistic_data) # 放入模型进行预测
_, pred = y_pred.topk(1, 1, largest=True, sorted=True)
Y_data_label = [index2label.get(i.tolist()) for i in label_data] # 转换为具体名称
pred_label = [
index2label.get(i.tolist()) for i in pred.view(-1).cpu().detach()
]
logger.info('Alpha:{}'.format(alpha))
display_model_performance_metrics(true_labels=Y_data_label,
predicted_labels=pred_label,
classes=label_list)
logger.info('Finished! (*￣︶￣)')