# 第18届中国大学生计算机设计大赛

# 漏洞分阶——NLP 文本大模型漏洞等级研判系统

**团队人员：王伊舟 李伟**

## 1、数据分析

利用聚类和热力图来分析赛事方提供的数据，帮助我们探寻解决问题的方式。

聚类采用降维和标准化处理后的数据，随后进行K-Means算法进行训练建模并根据误差平方和准则、轮廓系数、卡林斯基-哈拉巴斯指数作为聚类簇数的确定依据。再结合热力图判断数据集中各部分之间的影响关系，从中得到各部分数据之间的相关性。

经过数据深入挖掘分析，并参考得出的结论，我们决定将CVSS评分预测研判分解成八个子问题，每个子问题对应着CVSS Vector String中的各个指标，通过预测每项指标的结果，最终合成预测的CVSS Vector String，再通过其计算得出CVSS得分。



## 2、算法设计

我们使用多个NLP模型结合词干提取与词性还原、Shapley值同时配合5008个额外与漏洞信息相关的额外单词表设计了一个CVSS评分预判模型，并利用热力图显示的权重调入参数输入模型。该模型由八个子模型构成，每个子模型分别预测CVSS Vector String中的AV、AC、PR、UI、S、C、I、A八个指标，然后结合我们由热力图得出的权重，对需要预测CVSS评分数据集的漏洞描述进行标签预测（每个指标具有不同的标签，不同标签对CVSS得分有不同的影响），然后得出最终的通过预测结果计算预测的CVSS得分。

## 3、模型设计

使用NLP模型结合词干提取与词性还原、Shapley值同时配合5008个额外与漏洞信息相关的单词设计了一个CVSS评分预判模型，并利用热力图显示的权重调入参数输入模型。该模型由八个子模型构成，每个子模型分别预测CVSS Vector String中的AV、AC、PR、UI、S、C、I、A八个指标，然后结合我们由热力图得出的权重，对需要预测CVSS评分数据集的漏洞描述进行标签预测（每个指标具有不同的标签，不同标签对CVSS得分有不同的影响），然后得出最终的通过预测结果计算预测的CVSS得分。

## 4、文件内容介绍

acc文件夹存放模型准确度等数据相关截图；

com文件夹存放模型性能分析相关截图，包括Squared Errors分布图；

cvss文件夹包含一些计算cvss得分的工具包；

data文件夹存放比赛方提供的原始数据集；

dataset文件夹为训练所用数据集、验证集、测试集；

dataset_more文件夹为我们自行采集包含数据更多的数据集；

main文件夹存放本作品主要运行代码，包括decisionTree、K-Means、output。decisionTree文件夹存放决策树数据分析相关代码，K-Means文件夹存放聚类分析相关代码，output存放训练的模型和输出的结果；

nltk_data文件夹用来存放nltk工具包

vocab文件夹存放加入的漏洞特征单词表



## 5、部署运行方式

本作品编写/实验环境：

​        操作系统：Windows 11 23H2

​        实验GPU：RTX3070Ti Laptop 8G

​        CUDA：12.5

​        Torch：2.4.1

模型训练：

​        （1）下载安装对应版本的Python依赖包，可直接使用命令pip install -r requirements.txt；

​        （2）根据提示安装NLTK相关工具包；

​        （2）直接运行编写好的train.sh脚本，在Linux和Windows上均可运行。

测试：

​         （1）模型训练完成后直接运行test.sh脚本。

注：不同操作系统会存在路径识别问题，需要更改代码中的路径引用。

本作品完整代码已上传GitLink仓库和百度网盘。百度网盘包括已训练的模型，可直接使用。由于GitLink代码仓库对单一文件具有大小限制，故未上传我们自己训练的模型，需要自行训练，仅包含全部源代码。