/Learning-Based-Malware-Family-Classification-Methods

Primary LanguagePython

Learning-Based-PE-Malware-Family-Classification-Methods

本项目包含三类基于学习的PE恶意软件家族分类方法，分别是基于图像的、基于二进制的和基于反汇编的方法，还有一种检测恶意软件类间漂移的方法。

基于图像的家族分类方法

1. binary_to_gray.py（二进制文件转为灰度图）、bytes_to_gray.py（.bytes二进制资产文件转为灰度图）

2. gray_to_color.py

3. fromscratch表示不使用迁移学习，transfer表示用迁移学习，迁移ImageNet的训练参数，冻结卷积层对分类层进行重训练，aggressive表示迁移学习同时训练最后的特征表示层。



基于二进制的家族分类方法

⭐ 参考文献：‘Malware Classification Method Based on Word Vector of Bytes and Multilayer Perception’

1. 原始二进制文件预处理（1_process.py）

把 .byte文件中所有内容（除开头的行号以及未知字符"??"），都读取到一个一维数组中。将一维数组中连续出现5次及以上的"CC"和"00"删掉（这些连续的"CC"或"00“通常没有意义），进行文件预处理。并将处理好的内容输出到.txt文件中。

2. 将每个文件作为语料库，获得文件中256个字节的嵌入向量，将每个文件表示一个词向量升序矩阵（2_CBOW-gensim.py）

将一个.txt文件作为一个英文语料库，通过python的第三方库gensim获得文件中256个字节对应的嵌入向量。其中，word2vec使用CBOW模型，嵌入向量维度为256，窗口大小为5，迭代次数为5，使用负采样算法。之后，处理训练好的词向量矩阵，按字节升序排序，这里首先判断该文件包含的词典数是否为256（即该文件是否够256字节），是的话才升序排序；否则跳过该文件，并把文件名写到.txt文件中,稍后处理（共191左右个文件被跳过）。将每个文件的词向量升序矩阵保存为对应的.csv文件。

3. 可视化所有文件对应的灰度图，确认获得的词向量矩阵确实具有家族特征（3_log-to-image.py）

对获得的升序词向量升序进行处理，将每个元素标准化为[0, 255]间的整数，从而将向量矩阵转换为灰度图进行观察。

4. 使用MLP模型进行恶意软件家族分类（4_10-fold-MLP.py）

输入为未标准化的升序词向量矩阵（进行flatten，获得65536维），输出为所属家族。共10860个恶意文件，训练集：验证集：测试集 = 8:1:1，随机抽样，仅进行一次训练和测试，未用十倍交叉验证。


⭐ 参考文献：‘Malware Detection by Eating a Whole EXE’

1. 预处理数据（1_preprocess.py）

将二进制文件中的"??"和行首编号除去，将字节转换为整数。为了统一文件大小，用于后续的模型输入，统一将文件截断或填充为2MB大小。

2. 使用MalConv进行恶意软件家族分类（2_10-fold-Malconv-muti-GPU.py、2_10-fold-Malconv-single-GPU.py）

分别使用了多GPU训练（数据并行）和单个GPU训练。

3. 使用DeepConv进行恶意软件家族分类（3_10-fold-Deepconv.py）

4. 模型构建（DeepConv-model.py、MalConv-model.py）



基于反汇编的家族分类方法

⭐ 参考文献："Classifying Malware Represented as Control Flow Graphs using Deep Graph Convolutional Neural Network"

1. 提取.asm文件中的ACFG图，即顶点属性特征矩阵和邻接矩阵（主文件acfg_pipeline.py，调用了cfg_builder.py、dp.utils.py、python23_common.py）

在主文件中只用processTrainSet()函数即可，修改pathPrefix（数据集.asm文件存放路径）和labelPath（trainLabels.csv存放路径）。最后将提取的所有.asm文件的ACFG都放在 MSACFG.txt 文件中（2.84GB），和数据集放在一起。

2. 在pytorch_structure2vec-master/s2v_lib/ 下，输入make命令，创建build文件夹

3. 运行脚本cv_dgcnn.sh，该脚本关联cross_valid.py

脚本可以设置ACFG存放文件、使用哪个GPU、用到的参数文件、训练路径、结果输出文件等。注意，在用脚本载入训练好的模型进行测试时，需要改动参数norm_path=norm，对数据进行归一化处理。

4. 输出结果：

  • 每次交叉验证最后一轮epoch的模型被存在 msacfg_models；
  • 每次交叉验证训练集和验证集的混淆矩阵、每个epoch的实验值和MSACFG.txt存放在一起；
  • 十倍交叉验证的平均结果存放在maldefender目录下 （MSACFGGpu1Run0.csv）。

5. msacfg.hp设置实验的具体参数：

  • 学习率1r = 0.0001 (当连续3个epoch loss不减少，则将学习率 * 0.1)
  • batchsize = 10 （论文报告的最优batchsize值）
  • cvFold = 10 （十倍交叉验证，论文中是五倍）
  • numEpochs = 10

⭐ 参考文献："Modeling Malware as a Language"

1. 将文件按所属的家族存放在各自的文件夹 （collect_data_k_nn.sh）

malware_by_class存放按家族存放的数据集。

2. 将所有.asm文件进行parse，有两种方法（parse_data_k_nn.sh）

  • slow method 提取的是操作码序列，即design-1，放于parsed_malware_slow_method文件夹下；
  • fast method 提取的是映射好的整条汇编代码，即design-2，放于parsed_malware_fast_method文件夹下。
3. run_tests_gensim_k_nn.sh
  1. preprocess_parsed_malwares.sh：remove掉parsed_malware中的坏文件，Removing Bad Parsed Malware，文件数由10868->10767。
  2. test_gensim_k_nn.py：用gensim生成每个文件的词典，将词表示为向量，得到每个文件的向量矩阵。计算每两个文件间的文档距离，将所有距离其排序存放在 distances_win5_sgx.pickle 文件中。这一步十分耗费时间，对于design-1还快一些，处理10868个文件用7.8天，处理design-2很慢很慢，只计算了232个design-2文件的文档距离，大概用了30个小时左右，计算53592个文档距离。
  3. test_k_nn.py：加载sorted后的文档距离文件.pickle，对于k=1，计算每个家族的分类准确度，并将每个文件的分类结果写入 results_win5_sgx/ malware_lang_x_y，x表示class，y表示第几个文件。
  4. parse_results.sh：负责将每一类的分类结果综合放在一个文档中，results_win5_sgx/classification/classifications_1_nn_x ， x表示class。
  5. confusion_matrix.py：计算混淆矩阵的，将4中的结果统计成混淆矩阵。
  6. drift-7/new_multi_core.py：缓解上述2运行时时间开销极大的问题，采用多CPU核并行计算。

⭐ 参考文献："Malware identification using visualization images and deep learning"

1. 提取反汇编.asm文件中的Opcode序列（Process_Extracting_opcode.py）

2. 将Opcode序列转换为image（Process_Perfect_image_generator.py）

3. 对恶意软件进行家族分类（10-fold-MCSC.py）

4. 分类模型构建（Cmodel.py）



类间概念漂移检测方法的实现

1. 数据预处理（format文件夹）

因为训练概念漂移检测模型的数据集格式要求类似于MNIST数据集格式，所以用format文件夹中的代码进行转换，将原始数据集（PE恶意软件灰度图像）转换为MNIST格式的文件。

  • 首先，将每个家族按文件夹存放于test-images和training-images文件夹中，图片用灰度图。training-images中存放的是训练漂移检测模型的数据，test-images中存放待测家族，家族文件夹从0开始命名即可，即0,1,2,3,4，...。
  • 然后，batchea.meta.txt按序存放每个家族的名称，需自行修改。
  • 接着，resize.py将PE恶意软件转换的灰度图像进行resize，转换为256 plus 256维的图像。
  • 最后，用convert-images-to-mnist-format.py将图像数据集整合为MNIST格式数据集，得到4个和MNIST格式一样的文件：test-images-idx3-ubyte，test-labels-idx1-ubyte,train-images-idx3-ubyte,train-labels-idx1-ubyte。它们分别存放测试集和训练集（的图像和label）。

2. 训练基于三元组损失的卷积模型

将获得的4个文件复制到Triplet文件夹下使用，该文件夹下主要的代码是train.py和visualization.py，前者负责训练一个基于三元组损失的卷积模型对原始图像进行低维向量表示，后者用于将获得的低维向量表示进行2D降维可视化（tensorboard打开查看）。

3. 利用训练好的卷积模型，获得测试集的低维向量表示，进行基于MAD的距离度量，进行类间概念漂移检测。

代码Get_emb_label.py和main.py。