.
├── config # 配置文件夹
│ ├── cfgs.py # 配置文件,包括参数设置、算法选择等
│ └── __init__.py
├── data # 数据集文件夹
│ ├── img_gray_with_check # 恶意 Windows PE 文件灰度图像数据集(进行数据清洗)
│ ├── img_rgb # 恶意 Windows PE 文件 RGB 图像数据集
│ ├── malimg_paper_dataset_imgs # malimg 数据集
│ └── raw # 天融信恶意 Windows PE 文件原始数据集
├── models # 分类模型文件夹
│ ├── cnnmodels.py # CNN 模型
│ ├── __init__.py
│ └── swintmodels.py # Swin Transformer 模型
├── README.md # 源代码说明文件(本文件)
├── requirements.txt # 项目依赖
├── result # 验证结果文件夹
│ ├── out_gray # 验证实验 2 结果
│ ├── out_malimg # 验证实验 1 结果
│ └── out_rgb # 验证实验 2 结果
├── run.py # 程序入口文件
├── run.sh # 项目执行脚本文件
├── scripts # 运行脚本文件
│ ├── ensemble.py # 集成学习
│ ├── __init__.py
│ ├── traincnn.py # 训练及评估 CNN
│ ├── predict.py # 模型推理
│ └── trainswint.py # 训练 Swin Transformer
└── utils # 工具文件
├── batch_convert.py # 批量将恶意代码文件转换为图像
├── bin2img.py # 将单个恶意代码文件转换为图像,包含 4 种图像化算法
├── data_loader.py # 数据集导入
├── early_stopping.py # 早停类
├── __init__.py
├── plot.py # 结果可视化
├── split_sections.py # 分割文件的测试程序
└── utils.py # 其他工具函数
环境 | 说明 |
---|---|
CPU | 24 Intel(R) Xeon(R) CPU E5-2678 v3 @ 2.50GHz |
GPU | NVIDIA GeForce GTX 1080 Ti |
Python | 3.8.16 |
CUDA | 10.1.243 |
cudann | 7.6.4.38-1 |
操作系统 | CentOS Linux release 7.2.1511 |
依赖项 | 见 ./requirements.txt |
模型参数配置
vim ./config/cfg.py
恶意代码图像化
python3 -u ./utils/batch_convert.py
使用的图像化算法、输入输出路径等在 ./config/cfgs.py
中修改
数据处理
python3 -u ./utils/data_loader.py
训练集、测试集以及验证集的划分文件将保存在 ./out/data
中
CNN 模型训练
python3 -u ./run.py --model <MODEL_NAME> --mode train # MODEL_NAME 取自 ("resnet50" "vgg16" "googlenet" "densenet121" "shufflenet")
训练过程中的损失函数、精度等数据将保存在 ./out/data
中,模型权重将保存在 ./out/model
中
CNN 模型评估
python3 -u ./run.py --model <MODEL_NAME> --mode test # MODEL_NAME 取自 ("resnet50" "vgg16" "googlenet" "densenet121" "shufflenet")
评估过程中使用训练时保存在 ./out/model
中的模型,评估数据信息将保存在 ./out/data
中
集成学习模型评估
python3 -u ./run.py --model <MODEL_NAME> --mode ensemble # MODEL_NAME 取自 ("bayesian_combination" "stacking" "majority_voting" "distribution_summation")
集成学习过程中产生的数据信息将保存在 ./out/data
中,模型权重将保存在 ./out/model
中
模型推理
python3 -u ./run.py --model <MODEL_NAME> --mode predict --path <FILE_PATH> # MODEL_NAME 取自 ("resnet50" "vgg16" "googlenet" "densenet121" "shufflenet" "bayesian_combination" "stacking" "majority_voting" "distribution_summation")
python3 -u ./run.py --model <MODEL_NAME> --mode predict --path <IMAGE_PATH> --isimage # MODEL_NAME 取自 ("resnet50" "vgg16" "googlenet" "densenet121" "shufflenet" "bayesian_combination" "stacking" "majority_voting" "distribution_summation")
推理过程中使用训练时保存在 ./out/model
中的模型和 ./out/data
中的数据
推理示例:
模型的输入为 Windows PE 恶意代码文件或其按照指定算法转换而成的图像,模型输出为该恶意代码文件/图片的预测类别
本项目使用了 2 个数据集,分别为 malimg 恶意代码灰度图像数据集和天融信提供的 Windows PE 恶意代码文件数据集
数据集文件夹组织结构如下
.
├── dataset # 数据集文件夹
│ ├── Class1
│ │ ├── file1
│ │ ├── file2
│ │ └── ...
│ ├── Class2
│ │ ├── file3
│ │ └── ...
│ ├── Class3
│ │ └── ...
│ └── ...
malimg 恶意代码灰度图像数据集
来源:https://drive.google.com/file/d/1M83VzyIQj_kuE9XzhClGK5TZWh1T_pr-/view。规模和种类如下
Windows PE 恶意代码文件数据集
来源:天融信提供。规模和种类如下
类别 | 文件数量 | 转化后的图像数量 |
---|---|---|
Backdoor | 1000 | 897 |
Generic | 1000 | 1000 |
Trojan | 900 | 606 |
Variant | 1000 | 1000 |
Virus | 800 | 785 |
Worm | 1000 | 948 |
总计 | 5700 | 5236 |