- 战队名:Petrichor
- 比赛官网:https://www.datafountain.cn/competitions/507
近年来,各种勒索软件、木马、病毒、恶意挖矿程序等多种形式恶意软件不断涌现,恶意软件作者为逃避检测,在恶意软件组件中引入了多态性。虽然属于同一恶意软件“家族”的恶意文件具有相同形式的恶意行为,但由于编写者使用了各种策略不断修改和/或混淆,原本隶属同一家族的文件看起来像许多不同的文件,给检测识别的准确性带来了高度挑战。本赛题的任务是要我们通过分析给定的训练集中10个家族的恶意软件的特点,设计算法并构建AI模型,来对测试集中各个家族的恶意软件进行区分。
本文档描述了我们团队针对2021CCF BDCI数字安全公开赛“基于人工智能的恶意软件家族分类”这一赛题的解决方案及算法。我们的方案基于针对Ember 、TF-IDF 、Asm2Vec三个不同维度特征的特征工程(包括样本均衡、特征提取、特征融合和特征选择等),使用了梯度增强(XGBoost)、加权软投票和模型集成等人工智能策略。该方案最终取得了0.149394的多分类对数损失,在线上测试集排行榜中排名第二。
特征工程:PE文件的字节直方图、字节熵直方图、字符串信息特征。由于PE头被截取,我们通过在汇编文件中提取了节区、导入表和导出表信息;可读性字符串和操作码序列做词频-逆文件频率(TF-IDF);”精简“的操作码语义做asm2vec。
模型构建:只使用了单一XGBoost模型,集合加权软投票和模型集成。
我们提取的单特征可以分为三种类型:Ember 、TF-IDF 、Asm2Vec 。所有特征提取的脚本均在features.py
中实现,所有特征工程的脚本均在feature_engineering.py
中实现。
PE文件样本去除了所有头部信息,因此在开源的Ember静态特征提取方法中可以使用不需要头部字段的字节直方图 ByteHistogram()
,字节熵直方图 ByteEntropyHistogram()
,字符串信息特征 StringExtractor()
。
- 字节直方图:统计文件中0-255共计256个字节整数值的出现个数
- 字节熵直方图:近似文件熵与字节值的联合概率分布
- 滑动一个2048字节的窗口,步长为1024字节
- 计算每个2048字节窗口的熵
- 统计滑动窗口的(字节,熵值)对,最后转换成1x256维的特征向量
- 字符串信息特征:关于可打印字符串(0x20-0x7f,至少包含5个字符)的统计信息
- 字符串的数量、平均长度
- 字符串中字符的直方图和熵直方图
- 是否以具有特殊含义的字符串开头
C:\
- 路径http://或https://
- URLHKEY_
- 注册表项MZ
- 捆绑文件
由于PE文件中缺失了很多头部信息,所以我们选择了在汇编文件中提取来弥补。我们通过在汇编文件中提取了节区特征 SectionInfo()
、导入表特征 ImportsInfo()
和导出表特征 ExportsInfo()
。
- 节区特征
- 节区的统计特征:节区(section)个数、段(segment)个数、异常节区个数(节区名为空、大小为0)、可读可执行和可写的节区个数、是否存在调试段、重定位段、资源段或TLS段
- 节区大小:节区进行对齐处理前节的大小和实际在磁盘中所占的空间大小
- 入口节区信息:第一个可执行节区的节区名和相关属性
- 导入表特征
- 导入地址表中记录的动态链接库以及从各个动态链接库中导入的函数列表
- 导出表特征
- 导出函数的列表
我们提取了可读性字符串和操作码序列做词频-逆文件频率(TF-IDF),一个词或一段汇编代码在一个样本中出现次数越多, 同时在所有样本中出现次数越少, 越能够代表该样本。但是由于本赛题是10分类问题,而且样本极度不平衡,所以我们将样本数量锐减的家族(7,8,9)训练家样本不添加进词汇库。
我们在PE文件中提取了正则re.compile(b"[a-zA-Z]+")
匹配字符串,并对长于20的字符串进行拆分并保存了带有元音的字符串。
TF-IDF参数1:{'max_features':1000}
,单特征使用
TF-IDF参数2:{'max_features':300}
,特征融合
举例:
['tmEVPuuuW', 'uGEjDP', 'PEPVVVVVVPV', 'ExitProcess']
我们在汇编文件'Segment type: Pure code'
段中带有操作码的行中提取了操作码、第一个操作数的寄存器名称和汇编文件的注释内容。并且如果操作数为call
,我们提取了不为sub_, dword_, unkown_
开头的函数名。并在每个函数或者主程序之前添加;
符号做分词。
TF-IDF参数:{'stop_words' : [';'], 'ngram_range' : (1, 3), 'max_features':1000}
举例:
push ebp
mov ebp, esp
mov eax, 5170h
call __alloca_probe
cmp [ebp+arg_8], 0
loc_40102C: ; CODE XREF: sub_401000+23�j
mov edx, [ebp+arg_0]
test edx, edx
提取后:
['pushebp','movsbp','moveax','call__alloca_probe',';','cmpebp','movedx','testedx']
我们在汇编文件'Segment type: Pure code'
段中带有操作码的行中提取了”精简“的操作码语义(操作码,寄存器,注释内容),把每个函数抽象成一句话作为语料库,同样的我们排除了样本数量锐减的家族(7,8,9)训练家样本。
举例:
push [esp+lpProcName] ; lpProcName
push hModule ; hModule
call ds:GetProcAddress
test eax, eax
jnz short locret_1000101C
push 0FFFFFFFEh ; uExitCode
call ds:ExitProcess
sub_1000101F proc near ; CODE XREF: sub_100010CE+43�p
push offset ProcName ; "LpkTabbedTextOut"
call sub_10001000
push offset aLpkdllinitia_0 ; "LpkDllInitialize"
mov dword_10003244, eax
call sub_10001000
提取后:
"push esp lpProcName push hModule jnz push uExitCode call"
"push LpkTabbedTextOut call push LpkDllInitialize mov eax call"
我们取每个单特征去预测结果概率大于0.9的个数如下:
embers | section | imports | exports | ins | words | semantic | |
---|---|---|---|---|---|---|---|
5 | 548 | 569 | 568.0 | NaN | 574 | 551 | 572 |
3 | 441 | 955 | 447.0 | 89.0 | 396 | 225 | 626 |
6 | 651 | 869 | 878.0 | 555.0 | 933 | 867 | 912 |
0 | 1003 | 647 | 1434.0 | 1436.0 | 718 | 1395 | 1434 |
8 | 99 | 101 | 99.0 | NaN | 99 | 104 | 99 |
9 | 194 | 194 | NaN | NaN | 196 | 194 | 194 |
1 | 160 | 288 | 261.0 | NaN | 279 | 267 | 259 |
4 | 110 | 382 | 65.0 | 7.0 | 125 | 287 | 88 |
7 | 225 | 227 | 224.0 | NaN | 226 | 226 | 225 |
2 | 57 | 65 | 61.0 | NaN | 76 | 66 | 57 |
对TF-IDF的特征进行了类权重ExtraTreesClassifier
特征选择
classes_weights = class_weight.compute_sample_weight(
class_weight='balanced',
y=train_labels
)
selector=SelectFromModel(estimator=ExtraTreesClassifier(n_estimators=200)).fit(train_data, train_labels, sample_weight=classes_weights)
对ember section ins words
和 ember section ins semantic
进行了特征融合得到了更好的效果。
1 | embers | section | ins | words | semantic | em_sec_ins_words | em_sec_ins_sem |
---|---|---|---|---|---|---|---|
5 | 548 | 569 | 574 | 551 | 572 | 574 | 574 |
3 | 441 | 955 | 396 | 225 | 626 | 502 | 460 |
6 | 651 | 869 | 933 | 867 | 912 | 933 | 847 |
0 | 1003 | 647 | 718 | 1395 | 1434 | 720 | 716 |
8 | 99 | 101 | 99 | 104 | 99 | 103 | 99 |
9 | 194 | 194 | 196 | 194 | 194 | 194 | 194 |
1 | 160 | 288 | 279 | 267 | 259 | 250 | 265 |
4 | 110 | 382 | 125 | 287 | 88 | 262 | 287 |
7 | 225 | 227 | 226 | 226 | 225 | 226 | 226 |
2 | 57 | 65 | 76 | 66 | 57 | 67 | 66 |
在训练过程中我们讲训练好的模型去预测按家族分组的训练集并与真实标签计算logloss
,并将logloss
的负对数作为这个特征在单个家族的权重值,用于后续的软投票。其结果如下表,可以看到各个特征普遍对3,4类家族分类效果差。
类别权重表:
ember | section | imports | exports | words_1000 | semantic | ember_section_ins_words | ember_section_ins_semantic | |
---|---|---|---|---|---|---|---|---|
0 | 5.502961 | 5.816745 | 4.481392 | 4.716284 | 4.846862 | 6.715224 | 6.629952 | 6.682302 |
1 | 6.173242 | 5.848058 | 1.516962 | 0.000000 | 6.271196 | 6.649541 | 6.711672 | 6.772931 |
2 | 4.152618 | 4.169329 | 3.939775 | 0.000000 | 4.726480 | 4.513484 | 4.378887 | 4.390134 |
3 | 4.676251 | 4.614416 | 1.712744 | 0.000000 | 3.430825 | 5.280007 | 5.711640 | 5.827535 |
4 | 5.085630 | 5.577473 | 0.464740 | 0.000000 | 4.999519 | 5.628714 | 5.950920 | 6.037817 |
5 | 5.337548 | 5.995301 | 5.303306 | 0.000000 | 5.027418 | 5.963597 | 6.176870 | 6.170710 |
6 | 6.186197 | 5.891929 | 2.993289 | 0.000000 | 5.990634 | 6.508531 | 6.695083 | 6.677561 |
7 | 6.308145 | 5.481308 | 3.565423 | 0.000000 | 6.416205 | 6.354636 | 6.393423 | 6.421776 |
8 | 5.697321 | 5.733566 | 5.976484 | 0.000000 | 5.822914 | 5.848925 | 5.748133 | 5.784123 |
9 | 6.247548 | 6.030149 | 0.000000 | 0.000000 | 6.372997 | 6.346382 | 6.259929 | 6.250989 |
得到如下的权重后,我们首先通过加权软投票我们将多个模型的预测结果乘上权重相加,概率值最大的确定为样本类别。然后再从预测结果集合中选取其中预测为投票结果且概率值最大的作为该样本的最终预测结果。
我们选取了以下4组特征集合先做加权软投票的分别得到预测结果,再做求和平均得到最终的提交结果。
- ['ember', 'section', 'imports', 'exports']
- ['section', 'exports', 'words_1000', 'semantic']
- ['section', 'exports', 'words_1000', 'ember_section_ins_semantic']
- ['section', 'exports', 'ember_section_ins_words', 'ember_section_ins_semantic']
|-- code
|-- __init__.py
|-- features.py # 特征提取类函数
|-- feature_engineering.py # 特征工程函数
|-- model.py # 自定义模型
|-- train.py # 训练脚本
|-- predict.py # 预测脚本
|-- main.py # 训练脚本&预测脚本
|-- README.md # 解决方案及算法介绍
由于没有使用深度学习模型,故没有过多要求。
|-- malware_classification_petrichor
|-- codes # 代码文件
|-- data # 数据集映射
|-- train.sh # 执行模型训练过程
|-- predict.sh # 执行测试数据预测过程
numpy==1.21.2
gensim==4.1.0
tqdm==4.42.1
xgboost==0.90
pandas==1.3.3
joblib==0.14.1
scikit_learn==0.24.2