凌 祎，周安民，賈 鵬

（四川大學網(wǎng)絡(luò)空間安全學院，成都 610065）

0 引言

惡意軟件的規(guī)模和造成的經(jīng)濟損失逐年上漲，并且惡意軟件的產(chǎn)業(yè)化現(xiàn)象明顯，攻擊方法變化多樣，攻擊的目的性越來越明顯。近年來，惡意軟件檢測主要使用傳統(tǒng)動靜態(tài)檢測方法提取惡意軟件特征，包括利用控制流圖及字節(jié)流等信息。隨著深度學習的發(fā)展以及在惡意軟件檢測上的應用，研究者能夠?qū)⒏嗑S度的信息嵌入到代碼的特征向量中，能夠有效避免統(tǒng)計特征的局限性。受到近些年來自然語言處理發(fā)展的啟發(fā)，在處理匯編語言時可以將匯編文本看作自然文本來進行特征提取，而隨著經(jīng)典的word2vec逐步向BERT等基于大量文本的預訓練模型發(fā)展，基于匯編語言的惡意軟件檢測方法也可以采用類似的方法。雖然匯編語言可以看作自然語言來處理，但是其語法結(jié)構(gòu)和含義與實際的自然語言還是存在一定的區(qū)別，因此在預訓練部分需要自行進行訓練并重新定義新的語法結(jié)構(gòu)。

常見的基于深度學習的惡意軟件檢測方法，大多利用了統(tǒng)計特征，近幾年逐步也有利用自然語言處理的特征提取方法出現(xiàn)，依據(jù)自身的上下文及語義信息來生成特征向量。針對傳統(tǒng)自然語言處理和匯編語言的差別，本文提出了一種新的檢測方法，該方法利用自建數(shù)據(jù)集WUFCG 對基礎(chǔ)BERT 模型進行預訓練，一定程度上彌補了自然語言和匯編語言在嵌入過程中的差別；針對程序運行的特點，不再僅僅把程序基本塊相連而是在基本塊的基礎(chǔ)上再以函數(shù)作為圖節(jié)點進行嵌入。

1 檢測模型概述

本章節(jié)將介紹檢測模型的整體概況和關(guān)鍵技術(shù)細節(jié)。

1.1 控制流圖及函數(shù)調(diào)用圖

基本塊是程序運行的基本單元，由基本塊構(gòu)成的程序控制流圖能夠很準確地反映程序的特征，并且控制流圖天然形成圖結(jié)構(gòu)=（，），基本塊本身作為圖頂點（），控制流關(guān)系形成邊（）。但在程序?qū)嶋H運行時，由基本塊構(gòu)成的函數(shù)結(jié)構(gòu)也是客觀存在的，僅僅只考慮基本塊構(gòu)成的圖結(jié)構(gòu)會存在信息損失，本文在控制流圖（CFG）的基礎(chǔ)上引入了函數(shù)調(diào)用圖（FCG）結(jié)構(gòu)，使得分類效果更加精確。

1.2 特征提取及BERT預訓練

傳統(tǒng)的基于深度學習的惡意軟件檢測特征提取大多基于指令數(shù)量的統(tǒng)計特征，例如統(tǒng)計常量代碼、比較指令、調(diào)用指令、算數(shù)指令、MOV 指令等，近期也有一些基于經(jīng)典自然語言處理的特征提取方法取得了相較于統(tǒng)計特征更好的效果。該方法采用BERT對匯編文本進行特征提取，相較于經(jīng)典自然語言處理方法，針對惡意軟件環(huán)境利用惡意軟件數(shù)據(jù)集WUFCG 對BERT 進行預訓練，并對BIG2015 數(shù)據(jù)集進行文本嵌入，如圖1所示。

圖1 函數(shù)的嵌入框架

1.2.1 數(shù)據(jù)預處理

第一部分為預訓練階段的數(shù)據(jù)預處理，利用WUFCG 中的惡意樣本，首先將二進制文件進行逆向并對匯編文本預處理，在一個單獨的基本塊中，把基本塊中冗余信息剔除，每一行代碼內(nèi)部使用短橫線連接，每行代碼之間使用空格隔開，每一行代碼看作一個單詞，一個基本塊看作一個句子，在單個函數(shù)中，相連的基本塊另起一行，函數(shù)與函數(shù)之間以空行間隔，這樣單個函數(shù)則可以看作一段，如圖2所示。

圖2 文本預處理示意圖

第二部分BIG2015 數(shù)據(jù)集的預處理，利用數(shù)據(jù)集中已逆向完成的匯編文本（.asm文件）。

1.2.2 預訓練

匯編文本處理以后，再利用BERT預訓練模型訓練參數(shù)，在預訓練環(huán)節(jié)中，使用WUFCG 來訓練BERT。在后續(xù)的訓練中，主要利用BERT的向量模式對BIG2015 數(shù)據(jù)集中的匯編文本進行嵌入，即將基本塊匯編文本（單句）嵌入。BERT 模型本身不含任何參數(shù)信息，GOOGLE 團隊在完成BERT算法后利用自身計算資源使用大量自然語言文本來預訓練，而匯編語言的單詞遠沒有自然語言豐富，直接利用官方參數(shù)不夠精確。BERT 的預訓練部分核心為兩個任務，第一個任務為Masked Language Model，第二個任務為Next Sentence Prediction。Masked Language Model 旨在一個句子中隨機選中若干token 用于在預訓練中對其進行預測，被選中的詞語有80%的概率被替換為［MASK］，10%的概率被替換為其他token，10%的概率不被替換；Next Sentence Prediction（NSP）旨在預測兩個句子（基本塊）是否會相連，在本文的環(huán)境中即預測兩個基本塊是否會相連，在函數(shù)內(nèi)部結(jié)構(gòu)層面進行信息的提取，如圖3所示。

圖3 預訓練框架

1.3 訓練模型

本節(jié)將介紹整體的惡意軟件分類模型。在上一節(jié)中獲得了基本塊的嵌入形式，為了利用FCG 進行圖分類，還需要將整個函數(shù)中的基本塊圖池化，我們利用鄰接矩陣將函數(shù)內(nèi)部的信息進行兩次傳播，再通過平均池化獲得函數(shù)的嵌入表達形式即帶屬性的FCG，公式（1）如下。

在圖神經(jīng)網(wǎng)絡(luò)訓練階段，我們采用SAGPool來完成最終的圖分類。在圖神經(jīng)網(wǎng)絡(luò)中，圖的節(jié)點的特征信息和圖的結(jié)構(gòu)信息都會對最后的分類結(jié)果產(chǎn)生影響，僅僅只關(guān)注某一方面的信息是不準確的，經(jīng)典圖的池化方法Set2Set、SortPool、DiffPool 均在考慮結(jié)構(gòu)與節(jié)點信息方面有所欠缺，因此本文選用綜合考慮了節(jié)點特征信息和圖拓撲結(jié)構(gòu)信息的SAGPool 來對最終的FCG進行池化分類。

SAGPool 采用了Self-Attention 機制，改寫了GCN 經(jīng)典結(jié)構(gòu)，增加了自注意力參數(shù)，獲得用于 池 化 的 自 注 意 力 得 分∈R，公 式（2）如下。

其中idx 為前「」 個節(jié)點索引，為特征注意力mask，更多細節(jié)請參考原文。

SAGPool有兩種經(jīng)典模型，全局池化結(jié)構(gòu)和分層池化結(jié)構(gòu)，全局池化結(jié)構(gòu)相較于分層而言在小圖上表現(xiàn)更加優(yōu)異，而我們的樣本以小圖為主，因此選用全局池化結(jié)構(gòu)，圖神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)如圖4 所示，左為全局池化，右為分層池化，圖中Graph Convolution層為標準的圖卷積層。

圖4 SAGPool池化框架

2 實驗及分析

本節(jié)將評估我們的惡意軟件分類效果，實驗所使用深度學習框架為pytorch，服務器操作系統(tǒng)為Ubuntu16.04，CPU 為Intel Core E5-2630 2.60 GHz，GPU為GTX 2080Ti，顯存48 G。

2.1 數(shù)據(jù)集

實驗部分使用了兩個數(shù)據(jù)集，第一個數(shù)據(jù)集為自建數(shù)據(jù)集，是只含惡意樣本的WUFCG，共含惡意樣本33554 個，均為x86/i386 結(jié)構(gòu)，惡意樣本來自安全網(wǎng)站、安全廠商以及相關(guān)論壇，經(jīng)過預處理后形成了60G 訓練文本；第二個數(shù)據(jù)集為2015 年Kaggle 主辦的微軟惡意軟件分類比賽中所使用的惡意軟件數(shù)據(jù)集BIG2015，共包含9 個惡意軟件家族，共10868 個樣本，包含十六進制文件（不含PE 頭）和IDA PRO 生產(chǎn)的匯編文件，但匯編文件中包含一些無效的文件（因加殼導致.asm 文件為亂碼），數(shù)據(jù)集詳情如表1所示。

表1 BIG2015樣本

2.2 實驗結(jié)果及分析

實驗部分采用十折交叉驗證，主要關(guān)注準確率、精確率、召回率以及1 值，實驗結(jié)果如表2所示。

表2 實驗結(jié)果對比

由于設(shè)備原因，實驗部分BERT采用的參數(shù)為最小模式BERT-TINY。實驗結(jié)果表明，在參數(shù)量最小的BERT模式下，其分類結(jié)果已明顯優(yōu)于傳統(tǒng)方法。

3 結(jié)語

該方法基于BERT 和SAGPool 實現(xiàn)了對惡意軟件家族分類，并取得了理想的效果。該方法不僅可以用在惡意軟件家族分類上，在惡意軟件檢測、漏洞檢測方面亦可有一定的拓展，具有實際意義，在后續(xù)的工作中可以考慮修改SAGPool 中的部分圖卷積細節(jié)、增加BERT 參數(shù)量來提高最終的效果。