[2022] 基于NLP的威胁检出引擎

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

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

### v1.0

为了验证我的猜想,我使用了capstone作为反汇编引擎,将单个汇编向量化,结果如下:
![](https://key08.com/usr/uploads/2022/03/3817772337.png)
之后,我提取了10w个样本的.text段,并且将他们以一个单词一个单词的作为分割,也就是说,一个汇编作为一个单词,而一句话为一个程序。这样就可以被tf-idf转化后送入神经网络.
举个例子:
![](https://key08.com/usr/uploads/2022/03/1783623905.png)
部分代码如下

```cpp
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是因为以后方便移植.

```cpp
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)
```

效果如下:
![](https://key08.com/usr/uploads/2022/03/2142318123.png)
嗯...对整体模型毫无帮助...
所以问题出在哪? 哪里遇到了问题?

### 问题出在哪

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

1. 机器视角能用于做特征的非常少
2. 没有全局意义,比如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
>
>....

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

这是一个简单的例子:
![](https://key08.com/usr/uploads/2022/03/1655660906.png)
在这个例子中,我们就能获取一段具有实际意义的切片:

```cpp
mov
mov
mov
mov
mov
mov
call
test
jz
```

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

切片1:

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

切片2:

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

切片3:

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

现在 我们就具有了代码识别能力,如此反复,就会得到类似于如下的数据:
![](https://key08.com/usr/uploads/2022/03/670550250.png)

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%左右的识别率,但是因为样本数量严重不足+词库严重不足的情况下).

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

白帽Wiki

一只鸭子

白帽Wiki - 一个简单的wiki