[2022] 基于NLP的威胁检出引擎 huoji 启发检测,启发,机器学习,人工智能,NLP 2022-03-17 624 次浏览 0 次点赞 NLP引擎出自于我的一个简单的想法 - 既然人可以通过汇编看出软件是否是病毒 那么机器是否可以通过汇编看呢? 为了让机器能看得懂代码,我首先想到的是非常经典的NLP分类问题,所谓的NLP分类问题,就是训练的时候给一堆词语+词频,推理的时候先把句子中的词语分出来,再然后计算这些词语的词频,最后得出这条句子属于什么类别的结论 <!--more--> ### v1.0 为了验证我的猜想,我使用了capstone作为反汇编引擎,将单个汇编向量化,结果如下:  之后,我提取了10w个样本的.text段,并且将他们以一个单词一个单词的作为分割,也就是说,一个汇编作为一个单词,而一句话为一个程序。这样就可以被tf-idf转化后送入神经网络. 举个例子:  部分代码如下 ```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) ``` 效果如下:  嗯...对整体模型毫无帮助... 所以问题出在哪? 哪里遇到了问题? ### 问题出在哪 经过测试,直接送入汇编代码存在的问题包括不限于: 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 > >.... "激励"机制是为了划分出 "边界"后,在边界外寻找的一种机制.诺找不到/找到下一段影响,则回滚到边界.诺找到无影响/消除影响 则加入. 这是一个简单的例子:  在这个例子中,我们就能获取一段具有实际意义的切片: ```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)) { .... } ``` 现在 我们就具有了代码识别能力,如此反复,就会得到类似于如下的数据:  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%左右的识别率,但是因为样本数量严重不足+词库严重不足的情况下). 后续让我们来优化它.... 本文由 huoji 创作，采用 知识共享署名 3.0，可自由转载、引用，但需署名作者且注明文章出处。 点赞 0
加油