2021 CCF BDCI 数字安全公开赛 - 基于人工智能的恶意软件家族分类

• 战队名：Petrichor
• 比赛官网：https://www.datafountain.cn/competitions/507

赛题简介

近年来，各种勒索软件、木马、病毒、恶意挖矿程序等多种形式恶意软件不断涌现，恶意软件作者为逃避检测，在恶意软件组件中引入了多态性。虽然属于同一恶意软件“家族”的恶意文件具有相同形式的恶意行为，但由于编写者使用了各种策略不断修改和/或混淆，原本隶属同一家族的文件看起来像许多不同的文件，给检测识别的准确性带来了高度挑战。本赛题的任务是要我们通过分析给定的训练集中10个家族的恶意软件的特点，设计算法并构建AI模型，来对测试集中各个家族的恶意软件进行区分。

一、解题方案思路

本文档描述了我们团队针对2021CCF BDCI数字安全公开赛“基于人工智能的恶意软件家族分类”这一赛题的解决方案及算法。我们的方案基于针对Ember 、TF-IDF 、Asm2Vec三个不同维度特征的特征工程（包括样本均衡、特征提取、特征融合和特征选择等），使用了梯度增强（XGBoost）、加权软投票和模型集成等人工智能策略。该方案最终取得了0.149394的多分类对数损失，在线上测试集排行榜中排名第二。

特征工程：PE文件的字节直方图、字节熵直方图、字符串信息特征。由于PE头被截取，我们通过在汇编文件中提取了节区、导入表和导出表信息；可读性字符串和操作码序列做词频-逆文件频率（TF-IDF）；”精简“的操作码语义做asm2vec。

模型构建：只使用了单一XGBoost模型，集合加权软投票和模型集成。

二、特征工程

2.1 单特征提取

我们提取的单特征可以分为三种类型：Ember 、TF-IDF 、Asm2Vec 。所有特征提取的脚本均在features.py中实现，所有特征工程的脚本均在feature_engineering.py中实现。

2.1.1 Ember 特征

PE文件样本去除了所有头部信息，因此在开源的Ember静态特征提取方法中可以使用不需要头部字段的字节直方图 ByteHistogram(),字节熵直方图 ByteEntropyHistogram() ,字符串信息特征 StringExtractor()。

• 字节直方图：统计文件中0-255共计256个字节整数值的出现个数
• 字节熵直方图：近似文件熵与字节值的联合概率分布
  • 滑动一个2048字节的窗口，步长为1024字节
  • 计算每个2048字节窗口的熵
  • 统计滑动窗口的（字节，熵值）对，最后转换成1x256维的特征向量
• 字符串信息特征：关于可打印字符串（0x20-0x7f，至少包含5个字符）的统计信息
  • 字符串的数量、平均长度
  • 字符串中字符的直方图和熵直方图
  • 是否以具有特殊含义的字符串开头
    • C:\ - 路径
    • http://或https:// - URL
    • HKEY_ - 注册表项
    • MZ - 捆绑文件

由于PE文件中缺失了很多头部信息，所以我们选择了在汇编文件中提取来弥补。我们通过在汇编文件中提取了节区特征 SectionInfo()、导入表特征 ImportsInfo()和导出表特征 ExportsInfo()。

• 节区特征
  • 节区的统计特征：节区（section）个数、段（segment）个数、异常节区个数（节区名为空、大小为0）、可读可执行和可写的节区个数、是否存在调试段、重定位段、资源段或TLS段
  • 节区大小：节区进行对齐处理前节的大小和实际在磁盘中所占的空间大小
  • 入口节区信息：第一个可执行节区的节区名和相关属性
• 导入表特征
  • 导入地址表中记录的动态链接库以及从各个动态链接库中导入的函数列表
• 导出表特征
  • 导出函数的列表

2.1.2 TF-IDF 特征

我们提取了可读性字符串和操作码序列做词频-逆文件频率（TF-IDF），一个词或一段汇编代码在一个样本中出现次数越多, 同时在所有样本中出现次数越少, 越能够代表该样本。但是由于本赛题是10分类问题，而且样本极度不平衡，所以我们将样本数量锐减的家族（7，8，9）训练家样本不添加进词汇库。

（1）可读性字符串 StringExtractor.tfidf_features()

我们在PE文件中提取了正则re.compile(b"[a-zA-Z]+")匹配字符串，并对长于20的字符串进行拆分并保存了带有元音的字符串。

TF-IDF参数1：{'max_features':1000} ,单特征使用

TF-IDF参数2：{'max_features':300}，特征融合

举例：

['tmEVPuuuW', 'uGEjDP', 'PEPVVVVVVPV', 'ExitProcess']

（2）操作码序列 OpcodeInfo.tfidf_features()

我们在汇编文件'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']

2.1.3 Asm2Vec 特征

我们在汇编文件'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"

2.1.4 单特征预测结果图

我们取每个单特征去预测结果概率大于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

2.2 特征选择

对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)

2.3 特征融合

对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

三、模型构建

3.1 加权软投票

在训练过程中我们讲训练好的模型去预测按家族分组的训练集并与真实标签计算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

得到如下的权重后，我们首先通过加权软投票我们将多个模型的预测结果乘上权重相加，概率值最大的确定为样本类别。然后再从预测结果集合中选取其中预测为投票结果且概率值最大的作为该样本的最终预测结果。

3.2 多模型集成

我们选取了以下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 # 解决方案及算法介绍

五、建模算力与环境说明

5.1 建模算力

由于没有使用深度学习模型，故没有过多要求。

5.2 镜像内容

|-- malware_classification_petrichor
    |-- codes # 代码文件
    |-- data # 数据集映射
    |-- train.sh # 执行模型训练过程
    |-- predict.sh  # 执行测试数据预测过程

5.3 镜像预装包

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