源代码说明文件

源代码总体说明

.
├── 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

ysyszheng / Malware