Code Monkey home page Code Monkey logo

malware_classification_bdci's Introduction

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

赛题简介

近年来,各种勒索软件、木马、病毒、恶意挖矿程序等多种形式恶意软件不断涌现,恶意软件作者为逃避检测,在恶意软件组件中引入了多态性。虽然属于同一恶意软件“家族”的恶意文件具有相同形式的恶意行为,但由于编写者使用了各种策略不断修改和/或混淆,原本隶属同一家族的文件看起来像许多不同的文件,给检测识别的准确性带来了高度挑战。本赛题的任务是要我们通过分析给定的训练集中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 wordsember 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

malware_classification_bdci's People

Contributors

powerlzy avatar yuan2li avatar

Recommend Projects

  • React photo React

    A declarative, efficient, and flexible JavaScript library for building user interfaces.

  • Vue.js photo Vue.js

    🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.

  • Typescript photo Typescript

    TypeScript is a superset of JavaScript that compiles to clean JavaScript output.

  • TensorFlow photo TensorFlow

    An Open Source Machine Learning Framework for Everyone

  • Django photo Django

    The Web framework for perfectionists with deadlines.

  • D3 photo D3

    Bring data to life with SVG, Canvas and HTML. 📊📈🎉

Recommend Topics

  • javascript

    JavaScript (JS) is a lightweight interpreted programming language with first-class functions.

  • web

    Some thing interesting about web. New door for the world.

  • server

    A server is a program made to process requests and deliver data to clients.

  • Machine learning

    Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.

  • Game

    Some thing interesting about game, make everyone happy.

Recommend Org

  • Facebook photo Facebook

    We are working to build community through open source technology. NB: members must have two-factor auth.

  • Microsoft photo Microsoft

    Open source projects and samples from Microsoft.

  • Google photo Google

    Google ❤️ Open Source for everyone.

  • D3 photo D3

    Data-Driven Documents codes.