基于NLP的威胁检出引擎-先知社区

NLP引擎出自于我的一个简单的想法 - 既然人可以通过汇编看出软件是否是病毒那么机器是否可以通过汇编看呢?

为了让机器能看得懂代码,我首先想到的是非常经典的NLP分类问题,所谓的NLP分类问题,就是训练的时候给一堆词语+词频,推理的时候先把句子中的词语分出来,再然后计算这些词语的词频,最后得出这条句子属于什么类别的结论

v1.0

为了验证我的猜想,我使用了capstone作为反汇编引擎,将单个汇编向量化,结果如下:

之后,我提取了10w个样本的.text段,并且将他们以一个单词一个单词的作为分割,也就是说,一个汇编作为一个单词,而一句话为一个程序。这样就可以被tf-idf转化后送入神经网络.
举个例子:

部分代码如下

from sklearn.model_selection import train_test_split, GridSearchCV, KFold

tf_idf_transformer = TfidfTransformer()
vectorizer = CountVectorizer(max_features=5000)
y = np.array(list(csv_data.label.values))
x_data = np.array(list(csv_data.Data.values))
tf_idf = tf_idf_transformer.fit_transform(
    vectorizer.fit_transform(csv_data['OpCode']))
x_train_weight = tf_idf.toarray()

之后送到xgboost里面,之所以使用xgboost是因为以后方便移植.

model = XGBClassifier(**{"n_estimators":300,"max_depth":8})
model.fit(X_train, y_train, eval_set=[(X_train, y_train), (X_test, y_test)], early_stopping_rounds=5, verbose=True)

效果如下:

嗯...对整体模型毫无帮助...
所以问题出在哪? 哪里遇到了问题?

问题出在哪

经过测试,直接送入汇编代码存在的问题包括不限于:

机器视角能用于做特征的非常少
没有全局意义,比如push eax 只看push eax没有任何意义

为了解决这个问题,我们需要做语句切割,提取具体含义的代码块,而不是单个汇编

语句切割

我们的要做的这个NLP学习引擎不同于用VM的切割.并且我们不需要考虑去混淆、去虚拟化之类的.只是简单的分片.但是我们会面临一个巨大的问题: 压根不知道从哪开始切片,从哪结束切片.

因此我们需要引入一些简单的逻辑去解决他

为了分割词语,我们需要定义这个语句所影响的范围,如 push eax 所造成的影响是esp/rsp,而mov eax,ebp影响的是eax

我们简单的定义一下会造成影响的指令:

push

mov

xor

add

sub

...

我们简单定义一下会消除影响的指令:

pop

mov

sub

add

test

....

那么问题就简单了,所谓的分片就是寻找出有影响 <--> 消除影响的对应关系,此外我们不是追踪控制流,我们需要规定边界,一旦触及,则启用 "激励机制":

jmp

test

jnz

ret

....

"激励"机制是为了划分出 "边界"后,在边界外寻找的一种机制.诺找不到/找到下一段影响,则回滚到边界.诺找到无影响/消除影响则加入.

这是一个简单的例子:

在这个例子中,我们就能获取一段具有实际意义的切片:

mov
mov
mov
mov
mov
mov
call
test
jz

很明显,我们得到了类似于下面的代码的三个切片:

切片1:

a[0] = 1;
a[1] = 2;
a[2] = 3;
if (!function_unk_1()) {
    .....
}

切片2:

a[0] = 1;
a[1] = 2;
a[2] = 3;

切片3:

int x = X;
if (!function_unk_2(x)) {
    ....
}

现在我们就具有了代码识别能力,如此反复,就会得到类似于如下的数据:

now....

v2.0

把程序想象成一段话把我们提取的汇编想象成一个单词.现在我们就有把一句话拆分成单词的能力了:

现在只需要跟传统NLP一样,但值得注意的是,我们不再需要CountVectorizer的分词,直接送入tf-idf就行了.

这是我使用500个黑文件 500个白文件训练的结果:

Accuracy: 87.74%

请记住,NLP需要大量的样本训练,同时提词的时间复杂度很高,就1000个文件在我的垃圾电脑上跑了几个小时才出结果.

但是,在实际测试的时候,这个识别的效果挺让我震惊的:

100%|██████████| 1102/1102 [08:32<00:00,  2.15it/s]
扫描的文件总数 927 是病毒的文件总数 428

效果看起来还不错(虽然只有40%左右的识别率,但是因为样本数量严重不足+词库严重不足的情况下).

后续让我们来优化它....