机器学习在恶意软件检测中的应用

 [[188537]]

一、前言

机器学习是计算机科学的一个分支学科，目的在于赋予计算机从数据中学习的能力，使计算机能够有效利用当今互联网中存在的PB量级的数据，为人们在决策制定、任务执行方面提供帮助支持，这些工作对人们而言复杂度很高且耗时巨大。

恶意软件是企业和用户每天面临的紧迫威胁。无论是钓鱼邮件还是通过浏览器直接投放的漏洞利用工具，这些恶意软件可以与多种规避技术和其他安全漏洞相结合，将现有的防御系统远远抛在脑后。诸如Veil、Shelter等恶意软件框架已经被专业人士用于渗透测试中，取得了非常不错的效果。

今天我将向读者介绍机器学习如何在不使用特征值检测和行为分析方法前提下来检测恶意应用。

顺便提一句，像CylanceProtect、SentinelOne、Carbon Black之类的安全产品在特征值检测和行为分析方面做了很多工作，本文介绍的恶意软件检测框架不会涉及这些产品所使用的这两类技术。

二、机器学习简介

机器学习这个分支学科融合了数学中的多个领域，主要包括统计学、概率论、线性代数以及数学计算(如算法、数据处理、数值计算)。机器学习能够深入挖掘大数据价值，被广泛用于欺诈检测、垃圾邮件检测、电影推荐、饮食及产品购买推荐等各方面。亚马逊、Facebook以及Google等数百家公司也使用机器学习来改进他们的产品及服务。

机器学习主要方法有两种：有监督学习(supervised learning)和无监督学习(unsupervised learning)。有监督学习中，我们要处理的数据已事先打上标签，无监督学习则与之相反。两种方法都可以用于恶意软件检测，但我们主要关注第一种方法，因为我们的目标是对文件进行归类。

分类(classification)是有监督学习的一个子域，分类对象可以是二进制文件(恶意或非恶意软件)或其他类型对象(阿猫、阿狗、阿猪等等)，因此恶意软件检测属于二进制文件分类范畴。

机器学习的详细介绍不在本文范围内，你可以通过多种渠道了解详细信息，也可以查看附录中的资源来深入学习。

三、问题集

机器学习的工作流程包括定义问题、收集数据、整理数据(使数据符合训练要求)、使用算法处理数据。这一系列步骤需要消耗大量资源，因此对普通人而言，机器学习在具体实现上较为困难。这些步骤称之为机器学习的工作流程，也是机器学习所需的最少步骤。

对于本文设定的场景，我们首先需要定义工作流程：

1、首先，我们需要收集恶意软件样本，剔除大小小于10k的那些样本。样本数量越多越好。

2、其次，我们需要从样本中提取有意义的特征，这些特征也是我们研究的基础。所谓的特征指的就是能够描述对象的那些属性，比如，一栋房子的特征包括：房间数、房屋面积、房屋价格等。

3、提取特征后，我们需要对样本进行处理，构建样本数据集。数据集可以是一个数据库文件或一个CSV文件，以便于转化为数据向量，因为机器学习算法的计算对象是向量。

4、最后，我们需要一个衡量指标来评价二进制文件的分类结果。有多种指标可以用来衡量算法的性能，如ROC(Receiver Operating Characteristic，试者工作特征)、AUC(Area Under roc Curve，ROC曲线下面积)、混淆矩阵(Confusion Matrix)等。这里我们使用的是混淆矩阵指标，因为它能够反应结果的正确比率以及假阳性比率、假阴性比例。

四、收集样本以及特征提取

本文假设读者已经了解PE文件格式的相关知识，或者读者也可以先从这里学习基础知识。收集样本非常简单，你可以使用付费服务(如VirusTotal)或者使用这个链接中的样本源。

现在我们开始讨论建模问题。

为了让我们的算法能够从输入的数据中学习，我们需要清理数据，使之整洁且易于理解。本文中，我们使用12个特征来训练算法，这12个特征提取自样本文件，保存在CSV文件中。

(一)特征提取

我们使用pefile提取样本特征。首先是使用python下载pefile，命令如下：

pip install pefile 

工具准备完毕，在开始写代码前，我们先讨论一下我们到底需要提取哪些特征。对于一个PE文件来说，我们关心的主要是以下几个特征字段：

1、主映像版本(Major Image Version)：表示应用程序的主版本号。对于4.0版本的Excel而言，该值为4

2、IMAGE_DATA_DIRECTORY的虚拟地址以及大小

3、操作系统版本

4、导入地址表(Import Address Table)地址

5、资源区大小

6、区段个数

7、链接器版本

8、保留栈大小

9、DLL属性值

10、导出表大小和地址

为了使代码结构更为清晰，我们使用类对象来表示PE文件信息，类结构如下所示：

import os 
import pefile 
class PEFile: 
def __init__(self, filename): 
self.pe = pefile.PE(filename, fast_load=True) 
self.filename = filename 
self.DebugSize = self.pe.OPTIONAL_HEADER.DATA_DIRECTORY[6].Size 
self.DebugRVA =self.pe.OPTIONAL_HEADER.DATA_DIRECTORY[6].VirtualAddress 
self.ImageVersion = self.pe.OPTIONAL_HEADER.MajorImageVersion 
self.OSVersion = self.pe.OPTIONAL_HEADER.MajorOperatingSystemVersion 
self.ExportRVA = self.pe.OPTIONAL_HEADER.DATA_DIRECTORY[0].VirtualAddress 
self.ExportSize = self.pe.OPTIONAL_HEADER.DATA_DIRECTORY[0].Size 
self.IATRVA = self.pe.OPTIONAL_HEADER.DATA_DIRECTORY[12].VirtualAddress 
self.ResSize = self.pe.OPTIONAL_HEADER.DATA_DIRECTORY[2].Size 
self.LinkerVersion = self.pe.OPTIONAL_HEADER.MajorLinkerVersion 
self.NumberOfSections = self.pe.FILE_HEADER.NumberOfSections 
self.StackReserveSize =self.pe.OPTIONAL_HEADER.SizeOfStackReserve 
self.Dll =self.pe.OPTIONAL_HEADER.DllCharacteristics 
现在我们写个简单的函数，为每个PE文件构造一个字典，字典的键为特征字段，其值为特征值，这样每个样本都可以表示为一个python字典对象。如下所示： 
def Construct(self): 
sample = {} 
for attr, k in self.__dict__.iteritems(): 
if(attr != "pe"): 
sample[attr] = k 
return sample 

现在我们写个简单的函数，为每个PE文件构造一个字典，字典的键为特征字段，其值为特征值，这样每个样本都可以表示为一个python字典对象。如下所示：

def pe2vec(): 
    dataset = {} 
for subdir, dirs, files in os.walk(direct): 
for f in files: 
file_path = os.path.join(subdir, f) 
try: 
                pe = pedump.PEFile(file_path) 
                dataset[str(f)] = pe.Construct() 
except Exception as e: 
print e 
return dataset 
# now that we have a dictionary let's put it in a clean csv file 
def vec2csv(dataset): 
    df = pd.DataFrame(dataset) 
    infected = df.transpose()  # transpose to have the features as columns and samples as rows 
# utf-8 is prefered  
    infected.to_csv('dataset.csv', sep=',', encoding='utf-8') 

接下来我们准备处理这些数据。

(二)探索数据

这不是必要步骤，但可以让你对这些数据有直观上的理解。

import pandas as pd 
import numpy as np 
import matplotlib.pyplot as plt 
malicious = pd.read_csv("bucket-set.csv") 
clean = pd.read_csv("clean-set.csv") 
print "Clean Files Statistics" 
clean.describe() 
print "Malicious Files Statistics" 
malicious.describe() 

以下两个表格分别对应了正常程序和恶意文件的统计情况：

我们可以看到这两组数据集之间的差异，特别是前两个特征字段，差异更为明显。我们可以绘制一个图表，从直观上感受这些差异。

malicious['clean'] = 0 
clean['clean'] = 1 
import seaborn 
%matplotlib inline 
fig,ax = plt.subplots() 
x = malicious['IATRVA'] 
y = malicious['clean'] 
ax.scatter(x,y,color='r',label='Malicious') 
x1 = clean['IATRVA'] 
y1 = clean['clean'] 
ax.scatter(x1,y1,color='b',label='Cleanfiles') 
ax.legend(loc="right") 

图表如下：

从上图可知，恶意软件样本“聚类”程度较高，而正常文件样本稀疏分布在x轴上。接下来我们可以试着绘制其他特征的图表，以便全面了解这些样本数据。

分析“DebugRVA”特征：

%matplotlib inline 
fig,ax = plt.subplots() 
x = malicious['DebugRVA'] 
y = malicious['clean'] 
ax.scatter(x,y,color='r',label='Malicious') 
x1 = clean['DebugRVA'] 
y1 = clean['clean'] 
ax.scatter(x1,y1,color='b',label='Cleanfiles') 
ax.legend(loc="right") 

绘制的图表如下;

分析“ExportSize”特征：

%matplotlib inline 
fig,ax = plt.subplots() 
x = malicious['ExportSize'] 
y = malicious['clean'] 
ax.scatter(x,y,color='r',label='Malicious') 
x1 = clean['ExportSize'] 
y1 = clean['clean'] 
ax.scatter(x1,y1,color='b',label='Cleanfiles') 
ax.legend(loc="right") 

绘制的图表如下：

我们所绘制的图表越多，我们对数据的理解也越深，对数据的整体分布情况了解也越深。目前我们手上的数据集维度很低，那么问题来了，如果我们的数据集是高维度的，我们该如何处理?有很多技术可以降低数据集的维度，使“重要”特征更为突出。比如PCA和t-SNE算法可以将数据集绘制成三维甚至二维图像。

五、机器学习在恶意软件检测中的应用

前面我们已经做了足够多的统计工作，但在机器学习方面我们只做了一部分工作，如收集数据、清理及准备训练数据。在开始机器学习前，我们先要完成以下工作。

1、首先，我们需要将两部分数据集(Dataset)并为一个数据框(DataFrame)。

2、其次，我们需要数据框分为两部分，第一部分用于训练，第二部分用于测试。

3、接下来，我们将使用几个机器学习算法，看一下结果如何。

(一)数据集准备

import pandas as pd 
dataset = pd.read_csv('malware-dataset.csv') 
""" 
Add this points dataset holds our data 
Great let's split it into train/test and fix a random seed to keep our predictions constant 
""" 
import numpy as np 
from sklearn.model_selection import train_test_split 
from sklearn.metrics import confusion_matrix 
#let's import 4 algorithms we would like to test 
#neural networks 
from sklearn.preprocessing import StandardScaler 
from sklearn.neural_network import MLPClassifier 
#random forests 
from sklearn.ensemble import RandomForestClassifier 
""" 
Let's prepare our data 
""" 
state = np.random.randint(100) 
X = dataset.drop('clean',axis = 1) 
y = dataset['clean'] 
X = np.asarray(X) 
y = np.asarray(y) 
X = X[:,1:] 
X_train,X_test,y_train,y_test = train_test_split(X,y,test_size = 0.1,random_state=0) 

现在我们手上有了4个大型矩阵，其中X_train和y_train将用于训练不同的分类器，X_test用于标签预测，y_test用于指标衡量。事实上，我们将比较X_test和y_test的预测值，以便分析算法的具体实现。

(二)算法选择

首先来看看“Random Forests”(随机森林)算法，该算法是“决策树”(Decision Trees)算法的一种集成算法，核心思想是在训练期间内创建大量分类决策树，输出的分类即为样本分类的基础模型。随机森林算法在解决二进制文件分类问题上非常有效。

#let's start with random forests 
#we initiate the classifier 
clf1 = RandomForestClassifier() 
#training 
clf1.fit(X_train,y_train) 
#prediction labels for X_test 
y_pred=clf1.predict(X_test) 
#metrics evaluation 
""" 
tn = True Negative a correct prediction clean predicted as clean 
fp = False Positive a false alarm clean predicted as malicious 
tp = True Positive a correct prediction (malicious) 
fn = False Negative a malicious label predicted as clean 
""" 
tn, fp, fn, tp = confusion_matrix(y_test, y_pred).ravel() 
print "TN = ",tn 
print "TP = ",tp 
print "FP = ",fp 
print "FN = ",fn 

程序输出为：

TN = 697 
TP = 745 
FP = 6 
FN = 4 

根据处理结果，在没有进行参数微调和修改的情况下，我们只有6个假阳性和4个假阴性误判，这个结果相当不错。我们可以正确判断697个正常文件以及745个恶意软件，从结果上来看，我们的小型反病毒引擎效果还可以。

接下来我们试一下另一个分类器，我们建立一个简单的神经网络，看看它对随机分割的处理效果如何。

根据维基百科的词条解释：

多层感知器(multilayer perceptron，MLP)是一种前馈人工神经网络模型，它将输入数据集映射为一组适当的输出集。MLP由有向图中的多层节点组成，每层节点都与下一层节点完全相连。除了输入节点之外，每个节点都是具有非线性激活功能的神经元(或处理单元)。MLP使用了反向传播(back propagation)这种监督学习技术(supervised learning technique)来训练神经网络。MLP是标准线性感知器的修改版，可以用来区分不能线性分离的那些数据。

从上述定义我们可知，MLP是感知器的一种广义形式，也是深度学习方法的基本模型之一，可以用于处理广度和深度网络。

#our usual split 
X_train,X_test,y_train,y_test = train_test_split(X,y,test_size = 0.3,random_state=0) 
#This is a special process called feature engineering where we transform our data into the same scale for better predictions 
scaler = StandardScaler() 
scaler.fit(X_train) 
X_train = scaler.transform(X_train) 
X_test = scaler.transform(X_test) 
#Here we build a Multi Layer Perceptron of 12 Layers for 12 Features  you can use more if you want but it will turn into a complex zoo 
mlp = MLPClassifier(hidden_layer_sizes=(12,12,12,12,12,12)) 
#Training the MLP on our data 
mlp.fit(X_train,y_train) 
predictions = mlp.predict(X_test) 
#evaluating our classifier 
tn, fp, fn, tp = confusion_matrix(y_test,predictions).ravel() 
print "TN = ",tn 
print "TP = ",tp 
print "FP = ",fp 
print "FN = ",fn 

程序输出为：

TN = 695 
TP = 731 
FP = 8 
FN = 18 

看上去强大的神经网络并不能够识别样本数据集中的18个恶意软件(假阴性)，这是个很严重的问题，试想一下如果你的杀毒软件将勒索软件误判为正常程序，会对你造成什么影响?但不用过于悲观，因为我们这个神经网络还是非常原始的，实际上我们可以让它更为准确，但这已经超出了本文的讨论范畴。

六、总结

本文只是一篇入门文章，我想向读者表达的是，如果我们能够接受99%的识别率，那么恶意软件鉴别并不是一个难以解决的问题。当然，现实生活中，构建和部署机器学习是一件费时费事的工作，需要大量知识和大量数据。本文仅仅是机器学习和人工智能(AI)如何应用于恶意软件识别的一篇简单文章，希望能给读者提供学习知识的乐趣。