恶意软件分类模型——MalConv的实现

前言

文章原链接，《Malware Detection by Eating a Whole EXE》。文章中详细比较了之前的各种模型与MalConv模型的优劣。MalConv模型是对win32的PE文件进行分类。本文中实现采用了Google的开源框架Tensorflow。先来一张模型的图。

数据集的制作

数据集的制作采用Java，因为本人相对于Python更熟悉Java。

恶意文件

首先我手上有1.3G，共639个恶意文件样本。由于原文中是对大小2MB以下的恶意文件进行分类，这里需要对文件大小进行一下过滤。

/*
* 过滤大小不符合限制的恶意文件
* */
import java.io.File;

public class Main
{
    public static void main(String[] args)
    {
        File file = new File("samples/");
        int num = 0;
        for(File temp : file.listFiles())
        {
            if(temp.length() > 2000000)
            {
                temp.delete();
            }
        }
    }
}

注意此处过滤的阀值是2,000,000而不是2,097,152。为什么要采用2,000,000呢？在模型的卷积操作中卷积核的大小是500，如果阀值是2,097,152，那么就不能被500整除。过滤之后，恶意文件的个数为547。

正常文件

接下来是制作正常文件的数据集。

从windows C盘扫描符合大小的正常文件

我们从C盘扫描正常的文件，先来说说我遇到的几个坑。
首先我的电脑是Ubuntu和windows双系统，然后本篇是在Ubuntu下完成的。有一个大坑就是C盘的文件夹是存在软连接循环的。可以看到 /media/wuuuudle/1C30BB8E30BB6E00/Users/13648/AppData/Local 目录下的文件夹 Application Data 又软连接到了 /media/wuuuudle/1C30BB8E30BB6E00/Users/13648/AppData/Local ，形成了软连接循环。如果在扫描的时候不对软连接进行屏蔽，则会在此卡死。判断是否是软连接也很简单file.getAbsoluteFile().equals(file.getCanonicalFile())
第二个坑是，扫描整个C盘，一个扫描出来数万个符合大小的PE文件，但是恶意文件数据集的类型是32位，而扫描出来的正常PE文件是32位与64位均有。之后换了一个思路，不扫描整个C盘，而是扫描C盘下的 Program Files (x86) 目录，这样就只有32位的PE文件了。

/*
 * 扫描windows获得正常的windows PE文件
 * */

import java.io.*;
import java.nio.file.Files;

public class Scan_data
{
    static boolean isPe(File file)
    {
        if (file.length() > 2000000)//文件大小大于2M直接略过
            return false;
        try
        {
            FileInputStream in = new FileInputStream(file);
            char x1 = (char) in.read();
            char x2 = (char) in.read();
            return (x1 == 'M' && x2 == 'Z');
        }
        catch (Exception e)
        {
            e.printStackTrace();
        }
        return false;
    }

    static boolean isDir(File file)
    {
        if (file.isDirectory())
        {
            try
            {
                return (file.getAbsoluteFile().equals(file.getCanonicalFile()));
            }
            catch (Exception e)
            {
                e.printStackTrace();
            }
            return false;
        }
        return false;
    }

    static void saveFile(File file)
    {
        try
        {
            Files.copy(file.toPath(), new File("normal/" + file.getName()).toPath());
        }
        catch (IOException e)
        {
            e.printStackTrace();
        }
    }

    static void scandir(File file)
    {
        for (File temp : file.listFiles())
        {
            if (isDir(temp))
            {
                new Thread(() -> scandir(temp)).start();
            }
            else if (temp.isFile())
            {
                if (isPe(temp))
                {
                    saveFile(temp);
                }
            }
        }
    }

    public static void main(String[] args) throws Exception
    {
        String path = "/media/wuuuudle/1C30BB8E30BB6E00/Program Files (x86)";
        File file = new File(path);
        new Thread(() -> scandir(file)).start();
    }
}

为了加快扫描速度，在扫描时采用了多线程。

过滤后文件的随机抽取

在上一步中一共扫描出来6195个符合大小的PE文件，正常文件和恶意文件相差了一个数量级，所以肯定是要对扫描出的文件进行过滤，这里我们随机抽取了573个正常文件。

/*
 * 随机抽取文件
 * */
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.util.Random;
public class Random_Get
{
    static boolean saveFile(File file)
    {
        try
        {
            Files.copy(file.toPath(), new File("normal_f/" + file.getName()).toPath());
            return true;
        }
        catch (IOException e)
        {
            e.printStackTrace();
            return false;
        }
    }

    public static void main(String[] args)
    {
        String path = "normal";
        Random random = new Random();
        File[] files = new File(path).listFiles();
        for (int i = 0; i < 600; i++)
        {
            int index = random.nextInt(files.length);
            if (!saveFile(files[index])) i--;
        }
    }
}

随机混合

在这一步中将恶意文件与正常文件进行混合，打上标签（0为正常文件，1为恶意文件），分出900个数据用于训练，220个用于测试，并生成对应的xls。

搭建模型

神奇的梯度消失

起初，我是全部按照论文中给出的模型进行搭建，但是最后跑出分类的accuracy在50%上下浮动，接近于随机猜测。于是我就把网络中的softmax前一层的网络打印出来，发现全部为Nan，直接梯度消失了。我就去Google了一下这个模型，在NVIDIA的官网上发现了这个模型
模型和论文中的给模型有点区别，加了两层Relu的激活层。

模型代码

模型是使用Tensorflow实现的

from ReadData import getReadDataObject_test
from ReadData import getReadDataObject_train
import tensorflow as tf

input = tf.placeholder("int32", shape=[None, 2000000], name="input")
y_ = tf.placeholder("float", shape=[None, 2], name="y_")
training = tf.placeholder("bool", name="training")

#8-dembedding
embedding = tf.Variable(tf.random_normal([257, 8]), name="embedding")
x = tf.nn.embedding_lookup(embedding, input)

#slice
#slice.shape=[branch,2000000,4]
sliceA = tf.slice(x, [0, 0, 0], [-1, -1, 4])
sliceB = tf.slice(x, [0, 0, 4], [-1, -1, 4])

#1D Convolution
kernelA = tf.Variable(tf.random_normal([500, 4, 128]), name="kernelA")
kernelB = tf.Variable(tf.random_normal([500, 4, 128]), name="KernelB")

A = tf.nn.conv1d(sliceA, kernelA, stride=500, padding='VALID')
B = tf.nn.conv1d(sliceB, kernelB, stride=500, padding='VALID')
#Gating
#G0.shape=[branch*4000*128]
#G0 = A * tf.nn.sigmoid(B)
G0 = tf.nn.relu(A * tf.nn.sigmoid(B))

#MaxPooling
#P.shape=[branch*1*128]
P = tf.layers.max_pooling1d(G0, pool_size=4000, strides=1, padding='valid', data_format='channels_last')
#降维,P.shape=[branch*128]
P = tf.squeeze(P, [1])

#FC1
W_fc1 = tf.Variable(tf.random_normal([128, 128]), name="W_fc1")
P = tf.matmul(P, W_fc1)
P = tf.nn.relu(tf.matmul(P, W_fc1))
P = tf.layers.batch_normalization(P, training=True)

#FC2
W_fc2 = tf.Variable(tf.random_normal([128, 2]), name="W_fc2")
P1 = tf.matmul(P, W_fc2)
P1 = tf.layers.batch_normalization(P1, training=training)
#softmax
y_conv = tf.nn.softmax(P1, name="y_conv")

cross_entropy = -tf.reduce_sum(y_*tf.log(y_conv))
train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy)
correct_prediction = tf.equal(tf.argmax(y_conv, 1), tf.argmax(y_, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"), name="accuracy")

ReadData = getReadDataObject_train()
sess = tf.InteractiveSession()
sess.run(tf.initialize_all_variables())

for i in range((int)(len(ReadData.fp_list)/10*3)):
  batch = ReadData.next_batch(10)
  if i%2 == 0:
    train_accuracy = accuracy.eval(feed_dict={input:batch[0], y_: batch[1], training: False})
    print(sess.run(y_conv, feed_dict={input: batch[0], y_: batch[1], training: False}))
    print ("step %d, training accuracy %g"%(i, train_accuracy))
  train_step.run(feed_dict={input: batch[0], y_: batch[1], training: True})

#测试模型
ReadData_test = getReadDataObject_test()
accur = 0
for i in range((int)(len(ReadData_test.fp_list)/10)):
    batch = ReadData_test.next_batch(10)
    train_accuracy = accuracy.eval(feed_dict={input: batch[0], y_: batch[1], training: False})
    accur += train_accuracy * 10
    print("step %d, testing accuracy %g" % (i, train_accuracy))
batch = ReadData_test.next_batch(len(ReadData_test.fp_list) - (int)(len(ReadData_test.fp_list)/10)*10)
train_accuracy = accuracy.eval(feed_dict={input: batch[0], y_: batch[1], training: False})
accur += train_accuracy * len(ReadData_test.fp_list) - (int)(len(ReadData_test.fp_list)/10)*10
print("total testing accuracy %g" % (accur/len(ReadData_test.fp_list)))

sess.close()