機器學習在二進位代碼相似性分析中的應用

前言

本文介紹了3篇二進位代碼相似性分析的頂會技術，他們體現了二進位代碼相似性分析中一些最先進的思想。

第一篇是Genius技術，是在《基於神經網絡圖嵌入的跨平台二進位代碼相似性檢測》論文中作為對比技術介紹，它首次使用圖嵌入這個機器學習的概念去做二進位代碼相似性分析，它涉及到了聚類算法、圖比對、密碼本等技術，也為後兩篇論文打下了基礎。

第二篇是Gemini技術，它使用了更先進的Structure2vec算法計算圖嵌入，並融合Siamese網絡和神經網絡。

第三篇是騰訊科恩實驗室最近的頂會技術，它以Gemini的思路作為基礎，使用了更多的自然言語處理和深度學習算法去訓練程序控制流圖中基本塊語義信息、結構信息、順序信息。

最近，科恩實驗室頂會 《基於語義感知的神經網絡的二進位代碼相似度檢測》 讓我眼前一亮。因為當時我也正嘗試在二進位代碼分析中引入機器學習，正在試用的是隨機森林和adaboost算法（相較之下太low了），打算用它們做一個獨立於IDA的func_split，用來識別二進位代碼的函數。

然後，就去看了一下這個論文，看上去很高大上。但鬱悶的是，對提到的二進位相關知識都很熟悉，但對算法應用的原因和意義都很懵逼。就這樣，深度學習在二進位分析上的應用讓我產生了極大興趣，我決定好好研究下。

但是這個論文用的機器學習算法太多了，而且工作是基於另一篇論文的基礎，為了搞懂這個論文，我決定先去讀它的基線論文 《基於神經網絡圖嵌入的跨平台二進位代碼相似性檢測》 ，然後再來看看科恩頂會的創新和精妙之處 。

論文分析

基於神經網絡圖嵌入的跨平台二進位代碼相似性檢測

從論文標題說起：

圖嵌入

先解釋一個最重要的術語「圖嵌入」，圖嵌入的目的是降維，在保留圖特徵的前提下將圖經過非線性變換成一維向量。

相較於以前通過比對兩個二進位函數的CFG圖（函數的控制流程圖）的相似性去判斷代碼相似性（圖比對），現在提取CFG圖嵌入後對向量做相似性比較，其時間代價是不是更小，而且特徵更為抽象。

但是，利用CFG圖嵌入做二進位代碼的相似性檢測，前人已經做過了(Genius)，而且Genius在對CFG做嵌入之前，將CFG轉換為ACFG，即提取控制流程圖中平台無關的基本塊屬性和塊間屬性。

Genius使用二部圖匹配算法和密碼本作為基礎去計算ACFG的嵌入，具體做法是使用聚類算法訓練出一個具有代表性的ACFG的特徵向量集合，形成一個ACFG與特徵向量對應的密碼本，對於新的待計算ACFG，通過二部圖匹配算法與碼本中每個代表性的ACFG進行相似性比對，新ACFG的嵌入即為與其最相似的代表性ACFG的特徵向量。

Genius的圖嵌入生成方法時間代價較高，而且碼本質量取決於訓練數據規模。

Gemini（論文實現的原型）採用Structure2vec算法計算ACFG的圖嵌入，該算法從空間結構的相似性去定義相似度，通過評價函數整合節點和它n層鄰居的信息，將這些信息壓縮成一個有限維的非線性向量。

原產的Structure2vec算法並不能直接用於計算ACFG圖嵌入 ，Gemini使用了其變種，首先對ACFG每個節點使用變種的Structure2vec算法壓縮基本塊屬性和圖結構信息到一個特徵向量中，然後使用聚合函數將所有節點的特徵向量聚合成一個向量。

聽起來不錯吧， 它可以整合CFG中基本塊信息和基本塊間的結構信息 ， 並把它們映射成一個非線性向量的表示形式。

註：Structure2vec論文是2017年發表，Genius論文是2016年發表，也不怪Genius沒考慮到用Structure2vec算法去計算ACFG圖嵌入。

加入神經網絡的圖嵌入

到這裡，Gemini做二進位代碼的相似性分析基礎架構都有了（即基於Structure2vec算法生成2個二進位函數的ACFG圖嵌入，然後通過比較2個嵌入向量的相似性），神經網絡有什麼用？這時，神經網絡算法就該出場了，給它安排的活是訓練Structure2vec算法的參數，這時候有公式可能更直觀些。

如公式2所示，這是Gemini選定用來在Structure2vec模型中做非線性映射的模型，這個非線性映射的輸入xv是節點的基本塊屬性信息，u∈N(v)是節點的鄰接節點信息，經過映射可得到一個節點的特徵向量，而整個ACFG特徵向量是通過所有節點聚合得來。

其中，σ（·）是一個非線性變化， Gemini為了使非線性變化較為強大，就把σ（·）設計成n層全連接神經網絡 ；W1是一個d*p矩陣，d是xv（基本塊屬性）的長度，p是生成的嵌入長度，W1也是待訓練參數。

目前，總結來說就是，基礎是Structure2vec模型中加入神經網絡來訓練其非線性變化的模型，再用聚合函數將ACFG節點的特徵向量聚合成ACFG的圖嵌入，結合神經網絡的圖嵌入最終輸出的依然是ACFG圖嵌入（也可以稱為ACFG的特徵向量）。

Siamese網絡

在前面已經介紹過本文採用的二進位代碼相似比對的方式就是比較圖嵌入，那麼圖嵌入計算出來後，如何處理可以得到相似度的結果，答案就是Siamese網絡，其模型特點是在樣本量少的情況下識別類型。

如上圖所示，Siamese模型本身包括了目標降維模塊，正好上文的融入神經網絡的圖嵌入即可作為模型中的降維模塊。整體模型的輸入是兩個待比較二進位函數的ACFG，經過嵌入網絡得到降維的特徵向量u1和u2，然後計算二者餘弦距離，通過閾值判定結果是-1（非相似），還是1（相似）。當然，訓練過程都有真實的標籤，以便整個模型不斷調整參數。

跨平台

跨平台指的是什麼呢？即，二進位代碼相似性比較與二進位代碼所處平台、使用的編譯器和優化選項無關，訓練出的ACFG圖嵌入計算模型和使用的相似比對方法與以上條件無關。

又是如何做的呢？總結起來有兩點工作：

繼承了Genius提取二進位函數CFG中與平台無關性的屬性，6個基本塊屬性： 字符串常量、數字常量，傳遞指令數量，調用指令數量，運算指令數量，指令總數量； 1個基本塊間屬性： 基本塊子代數量。 需要注意的是，Gemini拋棄了Genius中使用的另一個基本塊間屬性「介數中心性」（betweenness），原因是該屬性提取較慢，與模型的最終效果影響也不大。

訓練數據採集的平台無關性，Gemini使用同一份源碼在不同系統平台、不同編譯器、不同編譯器優化選項的條件下生成大量帶標籤的樣本。 另外有一個樣本處理的細節是 不會將同一份源碼編譯後的二進位代碼分開在訓練集、驗證集和測試集。 個人認為這樣做有助於在訓練模型的過程中充分考慮到編譯的不同條件，以訓練出平台無關性模型。 在測試時，也可以用其它源碼生成的樣本來驗證未經訓練的數據是否符合模型。

總結標題

在介紹了論文中的圖嵌入、神經網絡、跨平台的概念及論文中的應用，那論文整體的思路也都明確了。總結下標題，目標是做跨平台二進位代碼的相似性檢查，方法是加入神經網絡的圖嵌入。

對論文的思考

對於二進位代碼，它的變化並不是無窮無盡的，因為它的產生是來源於機器規則，它看起來複雜難懂，但其內部必然藏著巨大的規律性 ，而深度學習的優勢就在於挖掘這種潛在的規律，我目前認為「深度學習是只要你有，我就一定能挖出來」。至於怎麼挖，用什麼樣的算法和模型，這需要人去建立二進位數據和模型的聯繫和關係，這篇論文是一個很好的結合，我相信未來，在這個方向上的繼往開來者會很多。

對於ACFG中的屬性，不同平台、編譯器、優化選項生成的二進位代碼， 有些屬性值（運算指令、call指令、指令總數，甚至常量值等）存在噪聲，它的成分並不是全部用來體現函數本身 ，有些僅是某個平台和編譯器或優化選項特有的，例如一些安全機制，棧隨機值保護、棧指針驗證、異常處理等，在不同平台和優化選項中，這些可能出現，也可能不出現。訓練模型時需要考慮這些，在模型真正用於實際時也需要考慮這些。

論文寫的非常好，從引言到原理闡述，其他方法的對比，對模型中特殊情況的考慮（例如超參問題），還有深度挖掘了這個項目的主線之外的價值點（例如再訓練問題），並做了相關的實驗。值得經常翻看，所以認真的把它翻譯成了中文，包括人家的答謝，英文水平有限，還望有看到的見諒。

註：論文翻譯可在博客的歷史文章中找到。

基於語義感知的神經網絡的二進位代碼相似度檢測

現在來看看科恩的論文，著重說說它與Gemini的不同之處吧。

Semantic-aware 模塊： BERT算法訓練基本塊語義與結構信息

Gemini沿用了它的基線論文Genius提取CFG屬性後（即ACFG）做圖嵌入,該篇論文認為這將損失大量的語義信息，故而在形成類似Gemini中ACFG的過程前，引入了NLP（自然語言處理），即使用BERT算法對每個基本塊序列進行訓練，語義訓練模型包含4個子任務，期望訓練出的模型包含以下4個維度的信息：

• Block的代碼語義
• Block的鄰接信息
• Block是否屬於CFG
• Block所屬平台、編譯器、及哪個優化選項

語義提取模型最終輸出類似於Genius和Gemini的ACFG圖，區別是該部分使用機器學習替代了人工提取（使用IDA腳本），而且包含基本塊屬性額外的信息， 輸出的基本塊特徵向量由機器學習得到 。

MPNN算法計算「ACFG」的圖嵌入

MPNN（消息傳遞網絡）是由Goog le科學家提出的一種模型，它本質上還是利用卷積神經網絡學習圖結構信息，只是形式框架上較為新穎，它通過收到鄰居的消息來更新自己的狀態，所以叫作消息傳遞網絡。（目前還沒有對卷積神經網絡和MPNN進行細緻學習，能理解的就是這麼多，可能有不準確的地方）。

不過，簡單的理解，該模型的作用可以類比Gemini的Structure2vec模型，他們的輸入都是CFG圖（其中基本塊屬性被量化成有限維特徵向量），輸出都是圖嵌入， 但這個圖嵌入不是最終的圖嵌入，它還要融入節點順序模型的訓練成果。

CNN算法訓練節點順序的特徵

使用CNN（卷積神經網絡）對CFG中節點相應順序或者說布局進行訓練也是該論文的一個重要創新部分。

由於前面已經包含了基本塊的屬性信息和CFG整體結構信息，這個模型就專注於訓練CFG的節點順序信息。

具體做法是，先把CFG節點布局抽象成01矩陣，然後對矩陣做卷積神經網絡，學習CFG節點順序相關特徵，最後輸出是一個特徵向量。

模型融合

該篇論文使用了三個模型從不同角度去訓練二進位代碼CFG中的信息，其中語義感知模型和結構感知模型是串聯的，順序模型可以與前面兩個並連，最後使用MLP（多層神經網絡）將兩條線上輸出的特徵向量映射為一個特徵向量。

融入三種模型後輸出的這個特徵向量才是和Gemini中基於神經網絡的Structure2vec模型訓練出的特徵向量是類似地位。

對論文的思考

科恩的論文對二進位本身特性考慮的更多，例如我在前邊論文思考中提到的平台、編譯器、優化選項產生的噪聲，他們在語義感知模型中加入了訓練Block屬於某個平台、編譯器和優化選項這個信息，不確定是否能消除這個噪聲，但至少考慮到了。

CFG還有一個很重要的信息，即節點的有向傳播（ 基本塊間的跳轉是有向的，不能逆代碼流程方向傳播 ），論文中的順序感知模型也許正好可以彌補這個方面的信息。

整個提取圖嵌入的流程在3個地方用了神經網絡算法，1個地方用了NLP算法，而且訓練無法並行，也不是端到端， 最長的鏈路為BERT -> MPNN -> MPL， 這個也許會在實用性上面打些折扣，不知道是否可以做些合適的取捨或訓練算法融合呢？

感悟

在漏洞挖掘領域中，當一門新技術、一個新工具、一些新代碼出現時，厲害的人往往可以收穫一大批CVE，靠的是日益累積的技術能力和漏洞挖掘領域中敏銳的嗅覺。我看頂級研究也是靠這兩樣，以2016年Genius論文為基礎，它使用了圖比對和密碼本去計算CFG圖嵌入。

到2017年Structure2vec論文出現時，Gemini能第一個想到用Structure2vec算法去替代圖比對和密碼本。BERT和MPNN也是2017年以後才出現的，科恩的論文能想到它們適應於解決ACFG人工生成問題和圖嵌入計算的問題。

一個有前途的基礎研究出來後，必定是遍地開花的結果，不過果實是留給準備充分的人！