前言
文章原链接,《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
78from 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()