最近在Kaggle上微軟發起了一個惡意代碼分類的比賽，并提供了超過500G的數據(解壓后)。有意思的是，取得第一名的隊伍三個人都不是搞安全出身的，所采用的方法與我們常見的方法存在很大不同，展現了機器學習在安全領域的巨大潛力。在仔細讀完他們的代碼和相關的論文后，我簡單的進行了一些總結與大家分享。
需要指出的是，(1)比賽的主題是惡意代碼的分類，不是病毒查殺(2)比賽采用的方法是純靜態分析的方法，不涉及行為分析等動態分析方法。
因此這不意味著這個方法能夠取代現有的方法，但是了解它能夠為安全研究人員提供一個嶄新的思路，至于能否在工業界應用仍待進一步研究。
80年代末期，隨著惡意代碼的誕生，反惡意代碼軟件（或稱反病毒軟件）隨之誕生。這個時期的惡意代碼所采用的技術較為簡單，這使得對應的檢測技術也較為容易，早期的反病毒軟件大都單一的采用特征匹配的方法，簡單的利用特征串完成檢測。隨著惡意代碼技術的發展，惡意代碼開始在傳播過程中進行變形以躲避查殺，此時同一個惡意代碼的變種數量急劇提升，形態較本體也發生了較大的變化，反病毒軟件已經很難提取出一段代碼作為惡意代碼的特征碼。在這種情況下，廣譜特征碼隨之誕生，廣譜特征碼將特征碼進行了分段，通過掩碼字節對需要進行比較的和不需要進行比較的區段進行劃分。然而無論是特征碼掃描還是廣譜特征，都需要在獲得惡意代碼樣本后，進行特征的提取，隨后才能進行檢測，這使得對惡意代碼的查殺具有一定的滯后性，始終走在惡意代碼的后面。為了針對變種病毒和未知病毒，啟發式掃描應運而生，啟發式掃描利用已有的經驗和知識對未知的二進制代碼進行檢測，這種技術抓住了惡意代碼具有普通二進制文件所不具有的惡意行為，例如非常規讀寫文件，終結自身，非常規切入零環等等。啟發式掃描的重點和難點在于如何對惡意代碼的惡意行為特征進行提取。特征碼掃描、查找廣譜特征、啟發式掃描，這三種查殺方式均沒有實際運行二進制文件，因此均可歸為惡意代碼靜態檢測的方法。隨著反惡意代碼技術的逐步發展，主動防御技術、云查殺技術已越來越多的被安全廠商使用，但惡意代碼靜態檢測的方法仍是效率最高，被運用最廣泛的惡意代碼查殺技術。
微軟提供的數據包括訓練集、測試集和訓練集的標注。其中每個惡意代碼樣本(去除了PE頭)包含兩個文件，一個是十六進制表示的.bytes文件，另一個是利用IDA反匯編工具生成的.asm文件。如下圖所示
Kaggle比賽中最重要的環節就是特征工程，特征的好壞直接決定了比賽成績。在這次Kaggle的比賽中冠軍隊伍選取了三個“黃金”特征：惡意代碼圖像、OpCode n-gram和Headers個數，其他一些特征包括ByteCode n-gram，指令頻數等。機器學習部分采用了隨機森林算法，并用到了xgboost和pypy加快訓練速度。
本文主要關注惡意代碼圖像和OpCode n-gram，以及隨機森林算法的應用。
這個概念最早是2011年由加利福尼亞大學的Nataraj和Karthikeyan在他們的論文 Malware Images: Visualization and Automatic Classification 中提出來的，思路非常新穎，把一個二進制文件以灰度圖的形式展現出來，利用圖像中的紋理特征對惡意代碼進行聚類。此后，有許多研究人員在這個思路基礎上進行了改進和探索。就目前發表的文章來看，惡意代碼圖像的形式并不固定，研究人員可根據實際情況進行調整和創新。
國內這方面的研究較少，去年在通信學報上有一篇《基于紋理指紋的惡意代碼變種檢測方法研究》，是由北京科技大學的韓曉光博士和北京啟明星辰研究院等合作發表的，目測也是僅有的一篇。
本節介紹最簡單的一種惡意代碼圖像繪制方法。對一個二進制文件，每個字節范圍在00~FF之間，剛好對應灰度圖0~255(0為黑色，255為白色)。將一個二進制文件轉換為一個矩陣(矩陣元素對應文件中的每一個字節，矩陣的大小可根據實際情況進行調整)，該矩陣又可以非常方便的轉換為一張灰度圖。
python代碼如下
#!python
import numpy
from PIL import Image
import binascii
def getMatrixfrom_bin(filename,width):
with open(filename, 'rb') as f:
content = f.read()
hexst = binascii.hexlify(content) #將二進制文件轉換為十六進制字符串
fh = numpy.array([int(hexst[i:i+2],16) for i in range(0, len(hexst), 2)]) #按字節分割
rn = len(fh)/width
fh = numpy.reshape(fh[:rn*width],(-1,width)) #根據設定的寬度生成矩陣
fh = numpy.uint8(fh)
return fh
filename = "your_bin_filename"
im = Image.fromarray(getMatrixfrom_bin(filename,512)) #轉換為圖像
im.save("your_img_filename.png")
利用該代碼生成的幾種病毒樣本圖像如下所示
用肉眼可看出，同一個家族的惡意代碼圖像在紋理上存在一定的相似性，不同的惡意代碼家族是有一定區別的。如何用計算機發現和提取這些紋理的相似特征用以分類呢？這就需要用到計算機視覺里的一些技術了。在Nataraj和Karthikeyan的論文中采用的是GIST特征，GIST特征常用于場景分類任務(如城市、森林、公路和海灘等)，用一個五維的感知維度來代表一個場景的主要內容(詳情請參考文獻[xx])。簡單來說，輸入圖像，輸出為對應的GIST描述符，如下圖所示
在matlab實現里面每個圖像的GIST描述符都用一個向量表示，最后用SVM完成分類訓練。
這已經遠遠超出了場景識別所能做的。不過，國外有學者利用一些類似前文生成那種不規則圖像來欺騙深度學習模型，如下圖所示
詳情請參考@王威廉老師的微博。當然，二者并沒有什么直接關聯，因為基于深度學習的圖像識別系統的訓練數據是一些有意義的圖像。但這是一個非常有意思的巧合，至于基于深度學習的圖像識別能否用于惡意代碼圖像的特征提取和分類，我認為是一個潛在的研究點，所能做的也不局限于此，如果有做深度學習的朋友可以伙同做安全的朋友一起研究交流。
n-gram是自然語言處理領域的概念，早期的語音識別技術和統計語言模型與它密不可分。n-gram基于一個簡單的假設，即認為一個詞出現的概率僅與它之前的n-1個詞有關，這個概率可從大量語料中統計得到。例如“吃”的后面出現“蘋果”或“披薩”的概率就會比“公路”的概率大(正常的語料中基本不會出現“吃公路”這種組合)，可以看出n-gram在一定程度上包含了部分語言特征。
將n-gram應用于惡意代碼識別的想法最早由Tony等人在2004年的論文N-gram-based Detection of New Malicious Code 中提出，不過他們的方法是基于ByteCode的。2008年Moskovitch等人的論文Unknown Malcode Detection Using OPCODE Representation 中提出利用OpCode代替ByteCode更加科學。
具體來說，一個二進制文件的OpCode n-gram如下圖所示
針對這次Kaggle比賽提供的數據，用python提取出其n-gram特征
#!python
import re
from collections import *
# 從.asm文件獲取Opcode序列
def getOpcodeSequence(filename):
opcode_seq = []
p = re.compile(r'\s([a-fA-F0-9]{2}\s)+\s*([a-z]+)')
with open(filename) as f:
for line in f:
if line.startswith(".text"):
m = re.findall(p,line)
if m:
opc = m[0][10]
if opc != "align":
opcode_seq.append(opc)
return opcode_seq
# 根據Opcode序列，統計對應的n-gram
def getOpcodeNgram(ops ,n = 3):
opngramlist = [tuple(ops[i:i+n]) for i in range(len(ops)-n)]
opngram = Counter(opngramlist)
return opngram
file = "train/0A32eTdBKayjCWhZqDOQ.asm"
ops = getOpcodeSequence(file)
opngram = getOpcodeNgram(ops)
print opngram
# output
# Counter({('mov', 'mov', 'mov'): 164, ('xor', 'test', 'setnle'): 155...
決策樹在我們日常生活中無處不在，在眾多機器學習的書籍中提到的一個例子(銀行預測客戶是否有能力償還貸款)如下圖所示
在這個在決策樹中，非葉子結點如“擁有房產”、“是否結婚”就是所謂的特征，它們是依靠我們的知識人工提取出來的特征。但如果對某個領域不了解，特征數量又較多時，人工提取特征的方法就不可行了，需要依靠算法來尋找合適的特征構造決策樹。
限于篇幅，決策樹的構造等過程本文不進行展開，網上相關資源非常多。(只要能夠充分理解熵和信息增益的概念，決策樹其實非常簡單)
隨機森林是一個非常強大的機器學習方法，顧名思義，它是用隨機的方式建立一個森林，森林里面有很多的決策樹組成，隨機森林的每一棵決策樹之間是沒有關聯的。在得到森林之后，當有一個新的輸入樣本進入的時候，就讓森林中的每一棵決策樹分別進行一下判斷，預測這個樣本應該屬于哪一類（對于分類算法），然后看看哪一類被選擇最多，就預測這個樣本為那一類。
隨機森林的思想與Adaboost里面的弱分類器組合成強分類器的思想類似，是一種“集體智慧”的體現。例如，屋子里面有n個人，每個人作出正確判斷的概率為p(p略高于0.5,這時每個人可視為一個弱分類器)，他們判斷的過程獨立互不影響，最終以多數人的判斷為準。這里我們不從數學上來推導，類比拋硬幣，對一枚均勻的硬幣，拋n次的結果中，正面和反面的次數是差不多的。而對一枚不均勻的硬幣，若出現正面的概率略大于反面，拋n次的結果中出現正面次數比反面次數多的概率就會很大。所以即使每個分類器的準確度不高，但是結合在一起時就可以變成一個強分類器。
如圖所示，將訓練數據有放回的抽樣出多個子集(即隨機選擇矩陣中的行)，當然在特征選擇上也可以進行隨機化(即隨機選擇矩陣中的列，圖中沒有體現出來)，分別在每個子集上生成對應的決策樹
決策過程如下圖所示(忽略畫風不一致的問題...)
但是在Kaggle比賽中冠軍隊伍采用的方法并不是從二進制文件生成的圖像，也不是從.bytes文件，竟然是從.asm文件生成的圖像，他們也沒有使用GIST特征，而是使用了前800個像素值作為特征，讓人非常費解。
我專門給隊伍里的Jiwei Liu同學發了一封郵件進行咨詢，他給我的答復是：GIST特征與其他特征綜合使用時影響整體效果，所以他們放棄了GIST特征，另外使用.asm文件生成圖像純屬意外發現...
至于為什么是前800個像素，他的解釋是通過反復交叉驗證得出的，具體原因不清楚。(在后文的分析中我會談談我的一些看法)
這部分的實現不復雜，他們選取n=4，在具體的特征選擇上通過計算信息增益選取每個分類與其他分類區分度最高的750個特征。
其他一些特征包括統計Headers，Bytecode n-gram(n=2,3,4)，分析指令流(將每個循環固定展開5次)來統計指令頻數，這些特征雖然不像前面提到的特征那么有效，但是確實可以在一定程度上提升最終成績。
冠軍隊伍的代碼是為了參加比賽而寫的，時間倉促，又是多人合作完成，導致組織結構很亂，且基本沒有注釋，可讀性較差。更重要的是自己想動手實踐一下，所以按照他們的思路寫了幾個簡單的程序，忽略了一些處理起來比較復雜或者難以理解的過程，代碼可以在我的github上下載
由于只是做一些簡單的試驗，不需要太多的數據(否則速度會非常慢)，我從微軟提供的訓練數據中抽取了大概1/10左右的訓練子集，其中從每個分類的中都隨機抽取了100個樣本(9個分類，每個樣本2個文件，共1800個文件)，這樣也不需要用到pypy和xgboost，只需要用到numpy、pandas、PIL和scikit-learn這些庫即可
友情提示：要進行這個實驗，首先確保有一個比較大的硬盤，推薦使用Linux系統。
這一步需要提前將完整訓練集解壓好，數量龐大，時間比較久。
#!python
import os
from random import *
import pandas as pd
import shutil
rs = Random()
# 讀取微軟提供的訓練集標注
trainlabels = pd.read_csv('trainLabels.csv')
fids = []
opd = pd.DataFrame()
for clabel in range (1,10):
# 篩選特定分類
mids = trainlabels[trainlabels.Class == clabel]
mids = mids.reset_index(drop=True)
# 在該分類下隨機抽取100個
rchoice = [rs.randint(0,len(mids)-1) for i in range(100)]
rids = [mids.loc[i].Id for i in rchoice]
fids.extend(rids)
opd = opd.append(mids.loc[rchoice])
opd = opd.reset_index(drop=True)
# 生成訓練子集標注
opd.to_csv('subtrainLabels.csv', encoding='utf-8', index=False)
# 將訓練子集拷貝出來(根據實際情況修改這個路徑)
sbase = 'yourpath/train/'
tbase = 'yourpath/subtrain/'
for fid in fids:
fnames = ['{0}.asm'.format(fid),'{0}.bytes'.format(fid)]
for fname in fnames:
cspath = sbase + fname
ctpath = tbase + fname
shutil.copy(cspath,ctpath)
本實驗中只用到了.asm文件，用到了.asm文件圖像特征(前1500個像素)和OpCode n-gram特征(本實驗取n=3,將總體出現頻數大于500次的3-gram作為特征保留)，實現代碼與前文基本一致，具體細節可參考完整代碼。
因為scikit-learn的存在，將機器學習算法應用到其他領域變得非常方便快捷。例如我們已經抽取了這些惡意代碼的OpCode n-gram特征("3gramfeature.csv")，利用scikit-learn即可快速訓練一個隨機森林
#!python
from sklearn.ensemble import RandomForestClassifier as RF
from sklearn import cross_validation
from sklearn.metrics import confusion_matrix
import pandas as pd
subtrainLabel = pd.read_csv('subtrainLabels.csv')
subtrainfeature = pd.read_csv("3gramfeature.csv")
subtrain = pd.merge(subtrainLabel,subtrainfeature,on='Id')
labels = subtrain.Class
subtrain.drop(["Class","Id"], axis=1, inplace=True)
subtrain = subtrain.as_matrix()
# 將訓練子集劃分為訓練集和測試集 其中測試集占40%
X_train, X_test, y_train, y_test = cross_validation.train_test_split(subtrain,labels,test_size=0.4)
# 構造隨機森林 其中包含500棵決策樹
srf = RF(n_estimators=500, n_jobs=-1)
srf.fit(X_train,y_train) # 訓練
print srf.score(X_test,y_test) # 測試
這里只對預測的準確度做一個簡單的評判。
在只應用.asm文件圖像特征(firstrandomforest.py)或者Opcode n-gram特征(secondrandomforest.py)的情況下，以及二者相結合的情況(combine.py)，準確率如下所示
由于隨機森林訓練的過程中存在一定的隨機性，因此每次結果不一定完全相同，但總的來說，二者結合的準確率通常要高出許多，基本可以達到98%以上的準確率，而且別忘了我們只用了不到1/10的數據
觀察.asm文件的格式，開頭部分其實是IDA生成的一些信息，如下圖所示
可以目測這個長度已經超出了800個像素(800個字節)，實際上這800個像素和反匯編代碼沒有關系！完全就是IDA產生的一些信息，更進一步的說，實際上冠軍隊伍的方法壓根與惡意代碼圖像沒有關系，實際上是用到了IDA產生的信息。
但是仔細觀察這些IDA信息，貌似長的都差不多，也沒有什么有價值的信息，為什么可以有效區分惡意軟件類型呢？
一個大膽的猜想是微軟提前將這些惡意代碼分好類，在調用IDA進行反匯編的時候是按照每個分類的順序進行的，由于未知的原因可能導致了這些IDA信息在不同分類上有細微差別，恰好能作為一個非常有效的特征！
https://www.kaggle.com/c/malware-classification/
https://github.com/bindog/ToyMalwareClassification
Malware Images: Visualization and Automatic Classification
Detecting unknown malicious code by applying classification techniques on OpCode patterns(墻裂推薦)
如何使用GIST+LIBLINEAR分類器提取CIFAR-10 dataset數據集中圖像特征，并用測試數據進行實驗
How Random Forest algorithm works
Supervised Classification with k-fold Cross Validation on a Multi Family Malware Dataset
我的博客 <http://bindog.github.io/blog/ > 我的郵箱 bindog 艾特 奧特路克 .com 歡迎大家與我交流~