基于人工智能的漏洞数据分类

前言

为及时跟踪国际信息安全趋势，需对国际公开的漏洞数据内容进行及时统计和梳理，例如CVE漏洞平台。CVE平台的漏洞信息包含有CVE编号、漏洞评分、漏洞描述等内容，其中漏洞描述含有对漏洞的利用条件、受影响的范围、漏洞可达到的效果（危害）等内容。为了更好地理解和持续研究，需将这些漏洞信息按照一定规则进行分类。而在此过程中，人工筛选分类效率较低，耗时耗力，利用人工智能，通过自然语言处理则可能更好地解决这一问题。本文提出一种多任务文本分类算法，通过共享BERT embedding，设计多任务的损失函数，端到端的解决漏洞数据分类问题。最后通过大量的实验及相关模型优化，最终获得本次大赛冠军。

大赛官网：基于人工智能的漏洞数据分类 竞赛 - DataFountain

一、赛题数据分析与建模

1.1、漏洞建模参数示意图

1.2、网络节点漏洞信息建模规则

我们需要把漏洞的描述信息，标注出以下几个标签：

• CVE-Number ： 即漏洞库中的漏洞cve编号，用来对漏洞详细信息进行收录、管理与识别等。
• Privilege-Required ：即漏洞利用所需的攻击者权限；攻击者只有先获取到利用该漏洞的前提权限，才能对该节点发动原子攻击，权限等级的枚举为：{admin/root、Nonprivileged、access、unknown}, admin/root表明攻击者获取到了最高权限，为单一标签；Nonprivileged表明攻击者可以以普通权限（非特权）的身份操作主机，access表明仅拥有节点的访问权限（网络可达）；unknown为从描述中获取不到该标签对应的内容，即根据描述得不到需要的权限信息。
• Attack-Vector ：攻击向量（介质）；攻击者是通过某种介质来利用该漏洞的，向量的枚举为{remote、Non-remote },remote表示可通过远程方式完成攻击，Non-remote表示不能通过远程的方式来完成攻击（即需要通过本地物理接触（cvss中的physical））。
• Impact ：漏洞利用的结果（完成漏洞利用获取到的权限），可分为三级，分别为（Impact-level1、Impact-level2、Impact-level3），Impact-level1的枚举为{ Privilege-Gained(RCE)、Information-disclosure、DoS、access、other}, Privilege-Gained(RCE)在Impact-level2下的枚举为{ admin/root、Nonprivileged、unknown}，（information-disclosure）在Impact-level2下的枚举为{local(credit)、other-target(credit)、other}, Impact-level3即local(credit)细化枚举为{user/root、Nonprivileged、unknown}、other-target(credit)枚举为{user/root、Nonprivileged、unknown}； Privilege-Gained(RCE)表示可获取到的权限：admin/root表示该漏洞利用可获取到该原子的最高权限，注意无论是root权限或者admin权限均为标签admin/root；Nonprivileged表示该漏洞利用可获取到该原子的普通非特权权限；unknown即为从描述中获取不到权限的信息，即为不清楚；Information-disclosure表示该漏洞为信息泄漏，而其中local(credit)表示该漏洞可获取到该原子的登录凭证；other-target(credit)表示该漏洞可获取到其他原子的登录凭证，other表示为该漏洞泄漏的是没有登录凭证类的其他信息；DoS表示该漏洞为拒绝服务漏洞；access表示可获取的网络的访问状态，（例如：通过攻击防火墙，只是修改策略，可造成部分本来不能访问的端口服务可正常访问），other即为其他不属于上述的影响。

1.3、赛题数据分析

根据我们在本次比赛中使用的算法，我们将类别标签映射成如图2：

​ 图2 类别标签映射

Attack-Vector类别标签分布，如下图3：

图3

Impact类别标签分布，如下图4：

图4

Privilege-Required类别标签分布，如下图：

图5

从图3、4、5可以得知，本次比赛的数据类别极不平衡，也对我们后续实验增加了许多困难。

二、模型架构

在本次比赛中，我们基于BERT[1]预训练语言模型，先后设计了两种方案进行实验。

2.1、方案1:基于BERT的3分类头的文本分类模型

如图6所示，我们对于每个大类，分别使用BERT+Linear分别对Impact、Attack-Vector及Privilege-Required三个大类进行分类。

图6 基于BERT的3分类头的文本分类模型

对于该模型，我们进行了相关实验，表1显示了在A榜上实验结果：

表1 A榜：三分类头实验结果

13分类 （trick）	4分类 （trick）	2分类 （trick）	合并性能
0.7 (freelb+ep3）	0.9111 (fgm+ep3)	0.984 (freelb+ep3)	0.571
0.69 (freelb+ep10）	0.92 (fgm+ep10）	0.981 （freelb+ep3）	0.5768
0.704 (freelb+ep10)	0.92 (fgm+ep10)	0.984 (freelb+ep3)	0.5632

从表中分析可以看出，将最优的分类头的最佳模型组合在一起得到的总体性能反而得不到最好结果。具体原因由下列示意图（7）分析演示得到，在各个模型取得性能都较高时，组合性能不一定是最好的：

图7

2.2、方案2:基于BERT的多任务端到端的文本分类模型

如图8，我们将BERT得embedding层共享，设计基于多任务的联合模型，端到端的训练分类模型，使各分类任务之间产生联系。则相应的损失函数为公式（1）：
$$Los{s}_{总}=los{s}_{cls2}+los{s}_{cls4}+los{s}_{cls13}$$

图8 基于BERT的多任务端到端的文本分类模型

结果：不使用ID LEAK的情况下，A榜：0.6035

2.3、模型优化

我们基于上述模型进行了大量实验进行优化，对于本次比赛，优化能提升的trick有：

• 拟合目标：使用F1-score替换acc作为评价指标；
• 类别不平衡：使用Focal Loss[2]函数；
• 早停机制[3]；
• 对抗训练：使用FreeLB[4]对抗训练，对bert的embedding进行扰动；
• 差分学习率；
• Dropout；
• 权重衰减。
• …

2.4、其他

在本次我们还尝试了基于lgb的特征工程方法，因为我们发现比赛数据存在数据leak，比如：ID存在leak，通过自己丰富的比赛经验和通过各种数据划分操作进行实验。初步确定这种情况如果B榜数据分布情况与A榜有差别的话容易导致性能暴跌，因此最后决定没有利用这一特征。

三、总结

本文基于BERT预训练语言模型，设计多任务学习目标，共享bert的embedding层，端到端的对漏洞数据进行分类，更加贴近本次任务的目标及现实场景，且训练速度及推理速度都具备优越性，更易于落地。

参考

[1] Devlin J, Chang M W, Lee K, et al. Bert: Pre-training of deep bidirectional transformers for language understanding[J]. arXiv preprint arXiv:1810.04805, 2018.

[2] Lin T Y, Goyal P, Girshick R, et al. Focal loss for dense object detection[C]//Proceedings of the IEEE international conference on computer vision. 2017: 2980-2988.

[3] Prechelt L. Early stopping-but when?[M]//Neural Networks: Tricks of the trade. Springer, Berlin, Heidelberg, 1998: 55-69.

[4] Zhu C, Cheng Y, Gan Z, et al. Freelb: Enhanced adversarial training for natural language understanding[J]. arXiv preprint arXiv:1909.11764, 2019.