機器學習在惡意軟件檢測中的應用
一、前言
機器學習是計算機科學的一個分支學科,目的在于賦予計算機從數(shù)據(jù)中學習的能力,使計算機能夠有效利用當今互聯(lián)網(wǎng)中存在的PB量級的數(shù)據(jù),為人們在決策制定、任務執(zhí)行方面提供幫助支持,這些工作對人們而言復雜度很高且耗時巨大。
惡意軟件是企業(yè)和用戶每天面臨的緊迫威脅。無論是釣魚郵件還是通過瀏覽器直接投放的漏洞利用工具,這些惡意軟件可以與多種規(guī)避技術(shù)和其他安全漏洞相結(jié)合,將現(xiàn)有的防御系統(tǒng)遠遠拋在腦后。諸如Veil、Shelter等惡意軟件框架已經(jīng)被專業(yè)人士用于滲透測試中,取得了非常不錯的效果。
今天我將向讀者介紹機器學習如何在不使用特征值檢測和行為分析方法前提下來檢測惡意應用。
順便提一句,像CylanceProtect、SentinelOne、Carbon Black之類的安全產(chǎn)品在特征值檢測和行為分析方面做了很多工作,本文介紹的惡意軟件檢測框架不會涉及這些產(chǎn)品所使用的這兩類技術(shù)。
二、機器學習簡介
機器學習這個分支學科融合了數(shù)學中的多個領(lǐng)域,主要包括統(tǒng)計學、概率論、線性代數(shù)以及數(shù)學計算(如算法、數(shù)據(jù)處理、數(shù)值計算)。機器學習能夠深入挖掘大數(shù)據(jù)價值,被廣泛用于欺詐檢測、垃圾郵件檢測、電影推薦、飲食及產(chǎn)品購買推薦等各方面。亞馬遜、Facebook以及Google等數(shù)百家公司也使用機器學習來改進他們的產(chǎn)品及服務。
機器學習主要方法有兩種:有監(jiān)督學習(supervised learning)和無監(jiān)督學習(unsupervised learning)。有監(jiān)督學習中,我們要處理的數(shù)據(jù)已事先打上標簽,無監(jiān)督學習則與之相反。兩種方法都可以用于惡意軟件檢測,但我們主要關(guān)注第一種方法,因為我們的目標是對文件進行歸類。
分類(classification)是有監(jiān)督學習的一個子域,分類對象可以是二進制文件(惡意或非惡意軟件)或其他類型對象(阿貓、阿狗、阿豬等等),因此惡意軟件檢測屬于二進制文件分類范疇。
機器學習的詳細介紹不在本文范圍內(nèi),你可以通過多種渠道了解詳細信息,也可以查看附錄中的資源來深入學習。
三、問題集
機器學習的工作流程包括定義問題、收集數(shù)據(jù)、整理數(shù)據(jù)(使數(shù)據(jù)符合訓練要求)、使用算法處理數(shù)據(jù)。這一系列步驟需要消耗大量資源,因此對普通人而言,機器學習在具體實現(xiàn)上較為困難。這些步驟稱之為機器學習的工作流程,也是機器學習所需的最少步驟。
對于本文設(shè)定的場景,我們首先需要定義工作流程:
1、首先,我們需要收集惡意軟件樣本,剔除大小小于10k的那些樣本。樣本數(shù)量越多越好。
2、其次,我們需要從樣本中提取有意義的特征,這些特征也是我們研究的基礎(chǔ)。所謂的特征指的就是能夠描述對象的那些屬性,比如,一棟房子的特征包括:房間數(shù)、房屋面積、房屋價格等。
3、提取特征后,我們需要對樣本進行處理,構(gòu)建樣本數(shù)據(jù)集。數(shù)據(jù)集可以是一個數(shù)據(jù)庫文件或一個CSV文件,以便于轉(zhuǎn)化為數(shù)據(jù)向量,因為機器學習算法的計算對象是向量。
4、最后,我們需要一個衡量指標來評價二進制文件的分類結(jié)果。有多種指標可以用來衡量算法的性能,如ROC(Receiver Operating Characteristic,試者工作特征)、AUC(Area Under roc Curve,ROC曲線下面積)、混淆矩陣(Confusion Matrix)等。這里我們使用的是混淆矩陣指標,因為它能夠反應結(jié)果的正確比率以及假陽性比率、假陰性比例。
四、收集樣本以及特征提取
本文假設(shè)讀者已經(jīng)了解PE文件格式的相關(guān)知識,或者讀者也可以先從這里學習基礎(chǔ)知識。收集樣本非常簡單,你可以使用付費服務(如VirusTotal)或者使用這個鏈接中的樣本源。
現(xiàn)在我們開始討論建模問題。
為了讓我們的算法能夠從輸入的數(shù)據(jù)中學習,我們需要清理數(shù)據(jù),使之整潔且易于理解。本文中,我們使用12個特征來訓練算法,這12個特征提取自樣本文件,保存在CSV文件中。
(一)特征提取
我們使用pefile提取樣本特征。首先是使用python下載pefile,命令如下:
- pip install pefile
工具準備完畢,在開始寫代碼前,我們先討論一下我們到底需要提取哪些特征。對于一個PE文件來說,我們關(guān)心的主要是以下幾個特征字段:
1、主映像版本(Major Image Version):表示應用程序的主版本號。對于4.0版本的Excel而言,該值為4
2、IMAGE_DATA_DIRECTORY的虛擬地址以及大小
3、操作系統(tǒng)版本
4、導入地址表(Import Address Table)地址
5、資源區(qū)大小
6、區(qū)段個數(shù)
7、鏈接器版本
8、保留棧大小
9、DLL屬性值
10、導出表大小和地址
為了使代碼結(jié)構(gòu)更為清晰,我們使用類對象來表示PE文件信息,類結(jié)構(gòu)如下所示:
- 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
- 現(xiàn)在我們寫個簡單的函數(shù),為每個PE文件構(gòu)造一個字典,字典的鍵為特征字段,其值為特征值,這樣每個樣本都可以表示為一個python字典對象。如下所示:
- def Construct(self):
- sample = {}
- for attr, k in self.__dict__.iteritems():
- if(attr != "pe"):
- sample[attr] = k
- return sample
現(xiàn)在我們寫個簡單的函數(shù),為每個PE文件構(gòu)造一個字典,字典的鍵為特征字段,其值為特征值,這樣每個樣本都可以表示為一個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')
接下來我們準備處理這些數(shù)據(jù)。
(二)探索數(shù)據(jù)
這不是必要步驟,但可以讓你對這些數(shù)據(jù)有直觀上的理解。
- 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()
以下兩個表格分別對應了正常程序和惡意文件的統(tǒng)計情況:
我們可以看到這兩組數(shù)據(jù)集之間的差異,特別是前兩個特征字段,差異更為明顯。我們可以繪制一個圖表,從直觀上感受這些差異。
- 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軸上。接下來我們可以試著繪制其他特征的圖表,以便全面了解這些樣本數(shù)據(jù)。
分析“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")
繪制的圖表如下:

我們所繪制的圖表越多,我們對數(shù)據(jù)的理解也越深,對數(shù)據(jù)的整體分布情況了解也越深。目前我們手上的數(shù)據(jù)集維度很低,那么問題來了,如果我們的數(shù)據(jù)集是高維度的,我們該如何處理?有很多技術(shù)可以降低數(shù)據(jù)集的維度,使“重要”特征更為突出。比如PCA和t-SNE算法可以將數(shù)據(jù)集繪制成三維甚至二維圖像。
五、機器學習在惡意軟件檢測中的應用
前面我們已經(jīng)做了足夠多的統(tǒng)計工作,但在機器學習方面我們只做了一部分工作,如收集數(shù)據(jù)、清理及準備訓練數(shù)據(jù)。在開始機器學習前,我們先要完成以下工作。
1、首先,我們需要將兩部分數(shù)據(jù)集(Dataset)并為一個數(shù)據(jù)框(DataFrame)。
2、其次,我們需要數(shù)據(jù)框分為兩部分,第一部分用于訓練,第二部分用于測試。
3、接下來,我們將使用幾個機器學習算法,看一下結(jié)果如何。
(一)數(shù)據(jù)集準備
- 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)
現(xiàn)在我們手上有了4個大型矩陣,其中X_train和y_train將用于訓練不同的分類器,X_test用于標簽預測,y_test用于指標衡量。事實上,我們將比較X_test和y_test的預測值,以便分析算法的具體實現(xiàn)。
(二)算法選擇
首先來看看“Random Forests”(隨機森林)算法,該算法是“決策樹”(Decision Trees)算法的一種集成算法,核心思想是在訓練期間內(nèi)創(chuàng)建大量分類決策樹,輸出的分類即為樣本分類的基礎(chǔ)模型。隨機森林算法在解決二進制文件分類問題上非常有效。
- #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
根據(jù)處理結(jié)果,在沒有進行參數(shù)微調(diào)和修改的情況下,我們只有6個假陽性和4個假陰性誤判,這個結(jié)果相當不錯。我們可以正確判斷697個正常文件以及745個惡意軟件,從結(jié)果上來看,我們的小型反病毒引擎效果還可以。
接下來我們試一下另一個分類器,我們建立一個簡單的神經(jīng)網(wǎng)絡(luò),看看它對隨機分割的處理效果如何。
根據(jù)維基百科的詞條解釋:
多層感知器(multilayer perceptron,MLP)是一種前饋人工神經(jīng)網(wǎng)絡(luò)模型,它將輸入數(shù)據(jù)集映射為一組適當?shù)妮敵黾LP由有向圖中的多層節(jié)點組成,每層節(jié)點都與下一層節(jié)點完全相連。除了輸入節(jié)點之外,每個節(jié)點都是具有非線性激活功能的神經(jīng)元(或處理單元)。MLP使用了反向傳播(back propagation)這種監(jiān)督學習技術(shù)(supervised learning technique)來訓練神經(jīng)網(wǎng)絡(luò)。MLP是標準線性感知器的修改版,可以用來區(qū)分不能線性分離的那些數(shù)據(jù)。
從上述定義我們可知,MLP是感知器的一種廣義形式,也是深度學習方法的基本模型之一,可以用于處理廣度和深度網(wǎng)絡(luò)。
- #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
看上去強大的神經(jīng)網(wǎng)絡(luò)并不能夠識別樣本數(shù)據(jù)集中的18個惡意軟件(假陰性),這是個很嚴重的問題,試想一下如果你的殺毒軟件將勒索軟件誤判為正常程序,會對你造成什么影響?但不用過于悲觀,因為我們這個神經(jīng)網(wǎng)絡(luò)還是非常原始的,實際上我們可以讓它更為準確,但這已經(jīng)超出了本文的討論范疇。
六、總結(jié)
本文只是一篇入門文章,我想向讀者表達的是,如果我們能夠接受99%的識別率,那么惡意軟件鑒別并不是一個難以解決的問題。當然,現(xiàn)實生活中,構(gòu)建和部署機器學習是一件費時費事的工作,需要大量知識和大量數(shù)據(jù)。本文僅僅是機器學習和人工智能(AI)如何應用于惡意軟件識別的一篇簡單文章,希望能給讀者提供學習知識的樂趣。