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

前言

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

数据集的制作

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

恶意文件

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
/*
* 过滤大小不符合限制的恶意文件
* */
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文件了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
/*
* 扫描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个正常文件。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
/*
* 随机抽取文件
* */
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实现的

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
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()

文章目录
  1. 1. 前言
  2. 2. 数据集的制作
    1. 2.1. 恶意文件
    2. 2.2. 正常文件
      1. 2.2.1. 从windows C盘扫描符合大小的正常文件
      2. 2.2.2. 过滤后文件的随机抽取
    3. 2.3. 随机混合
  3. 3. 搭建模型
    1. 3.1. 神奇的梯度消失
    2. 3.2. 模型代码