CN113591093A - 基于自注意力机制的工业软件漏洞检测方法 - Google Patents

Abstract

本发明公开了一种基于自注意力机制的工业软件漏洞检测方法，该方法采用自注意力机制对软件漏洞进行检测，检测过程分为代码预处理过程和漏洞检测模型训练测试过程；首先获取数据库漏洞程序，提取库API函数调用和程序片段，分为5类漏洞问题构建成工业软件漏洞的漏洞库，然后对漏洞库的关键点进行切片，将切片后的程序组装成行关联代码并进行数据处理，之后将行关联代码转换成对应的向量获得特征向量，并处理成相同长度大小；将特征向量与位置编码向量相加作为漏洞检测模型的输入，然后训练D_transformer神经网络，将训练好的模型通过测试样本验证模型的检测能力。本发明方法进一步提高了检测分类精度，降低了漏报率。

Description

基于自注意力机制的工业软件漏洞检测方法

技术领域

本发明涉及网络安全技术领域，特别涉及一种基于自注意力机制的工业软件漏洞检测方法。

背景技术

随着各行业软件开发技术的不断发展，工业软件也应运而生，工业软件作为工业互联网应用体系的重要组成部分，是为满足特定需求的工业应用软件，可以将人的工作从枯燥、重复的体力劳动中解放出来，专注于更具价值的知识创造性工作，并可低成本享受先进的软硬件技术，提升企业生产效率，通过规模化的复用，进一步提升行业的智能制造水平。在工业软件为工业技术人员提供便利的同时，工业软件漏洞问题也成为软件开发者们不可忽视的重要问题，在工业互联网环境中，一旦软件连接了系统，如果软件出现漏洞问题，机器和设备就可能会遭到攻击，生产过程会受到破坏和干扰，甚至停止。因此，越及时检测出软件中潜在的漏洞问题，生产过程中产生的财产和安全损失也越小。传统的静态检测方法的误报率较高，动态检测方法容易产生漏报并且比较耗时。近几年来随着深度学习技术的发展，深度学习在很多领域取得了惊人的成就，深度学习能够学习具有更复杂和抽象的高级特征表示，能够自动学习更可概括的潜在特征或表示，并且提供灵活性，允许针对不同的应用场景定制网络结构，利用深度学习技术从大量代码中学习软件漏洞模式成为了可能。

卷积神经网络在计算机视觉方面对图片能够较好地获得更细粒化的特征，用来学习结构化的空间数据，对于单词或句子，它只能捕捉单词的上下文含义，因此应用卷积神经网络只能学习上下文的简单代码语义。循环神经网络在自然语言处理方向取得了令人惊叹的成果，它用于处理顺序数据，尤其是循环神经网络的双向形式能够捕获序列的长期依赖关系，因此许多研究使用双向长短时记忆网络和门控递归单元结构来学习代码上下文依赖关系。常用的深度学习模型还应用了注意力机制，只选择一些关键的信息输入进行处理，即对输入序列中不同时刻分配不同的注意力，来提高神经网络的效率，这对于理解许多类型的漏洞的语义至关重要。

发明内容

本发明的目的是进一步提高软件漏洞检测的精度和效率，同时还能在网络环境发生变化时表现出一定的自适应能力，提高对未知漏洞的检测能力，建立一个更为完善和全面的检测模型，从而改善软件漏洞的检测效果。本发明提供一种基于自注意力机制的工业软件漏洞检测方法，旨在提高检测性能。

为了解决上述存在的技术问题，实现发明目的，本发明是通过以下技术方案实现的：一种基于自注意力机制的工业软件漏洞检测方法，该方法采用自注意力机制对工业软件漏洞进行检测，检测过程分为代码预处理过程和漏洞检测模型训练测试过程：

所述的代码预处理过程的步骤为：

S11、首先根据数据库漏洞程序，提取API函数调用和程序片段，将API函数分为5类漏洞问题构建成工业软件漏洞的漏洞库，以漏洞库的关键点作为程序切片的切入点，提取出代码中与关键点有关的参数、语句和表达式；S12、利用切片工具对漏洞库的关键点进行切片，获得切片后的程序，然后将切片后的程序组装成行关联代码，对行关联代码进行数据处理，并将行关联代码转换成对应的向量，获得特征向量；S13、将特征向量统一处理为相同长度的向量，当行关联代码向量长度不足设置的长度时，则在向量最后补0，当行关联代码向量长度超过设置的长度时，超过的部分删除；

所述的代码漏洞检测模型训练和测试过程为：

S21、训练过程采用批量抽取的方式从训练样本中每次抽取固定个数的样本，传入到采用自注意力机制的代码漏洞检测模型中，获得预测结果；S22、将预测结果与真实结果的误差通过反向传播及梯度下降算法进行参数更新，通过多次迭代训练模型；S23、测试过程使用训练过的模型进行测试，对比模型测试结果与实际结果是否相同，测试模型的检测能力。

优选地，所述数据库采用NIST的漏洞数据库的源代码，包含软件产品中漏洞的NVD和学术安全漏洞的SARD，随机选择80％的程序作为训练程序，20％的作为测试程序。

优选地，所述5类漏洞问题分别为缓冲区溢出问题、空指针引用问题、堆溢出问题、API误用和信息泄露问题。

优选地，所述切片包含前向切片和后向切片，其中前向切片对应于受相关参数影响的语句，而后向切片对应于可影响相关参数的语句，利用数据依赖图，来提取这两种切片。

优选地，行关联代码为语义上相互关联的代码行。

优选地，所述数据处理主要是对行关联代码添加标签，包含漏洞的标签设为1，不包含漏洞的标签设为0，对行关联代码做分词处理和去重处理以及针对数据样本不均衡问题为样本赋予不同的权重。

优选地，所述将行关联代码数据转换成对应的向量采用基于层序softmax的word2vec的跳字模型。

优选地，所述代码漏洞检测模型为采用自注意力机制的D_transformer模型，其网络结构采用encoder-decoder架构，内部的encoder和decoder的各包含一个自注意力机制层和一个前馈神经网络。

与现有技术相比，本发明具有以下有益效果：

(1)本发明方法采用自注意力机制，对于长句子，对每个输入项分配的权重取决于输入项之间的相互作用，不会存在长期依赖问题，它能帮助当前节点不仅仅只关注当前的词，从而能获取到上下文的语义；

(2)本发明采用多头注意力机制，提高了模型的拟合能力；

(3)本发明使用D_transformer模型，可以同时接受所有的向量作为输入，具有并行计算的优势，因此训练速度更快，使用较深的模型，较深的模型相比于较浅的模型可以收敛于较低的测试误差，所以该模型可以获得更高的精度，使得漏洞检测方法更好的满足工程实际要求。

附图说明

图1是本发明所提方法的流程图；

图2是本发明的代码预处理流程图；

图3是本发明代码漏洞检测模型训练和测试流程图；

图4是本发明D_transformer模型的网络结构图；

图5是本发明encoder和decoder的内部简版结构图；

图6是本发明D_transformer的详细结构图。

具体实施方式

下面将结合本发明实施例中附图，对本发明实施例中的技术方案进行清楚、完整地描述，显然，所描述的实施例仅仅是本发明的一部分实施例，而不是全部实施例。通常在此处附图中描述和展示出的本发明实施例的组件可以以各种不同的配置来布置和设计。因此，以下对在附图中提供的本发明实施例的详细描述并非旨在限制本发明要求保护的范围，而是仅仅表示本发明的选定实施例。基于本发明的实施例，本领域技术人员在没有做出创造性劳动的前提下所获得的其他所有实施例，都属于本发明保护的范围。

下面结合附图与具体实施方式对本发明做更为详细的说明：

本发明的硬件实验环境为:CPU Intel Core I7-9750H，2.60GHZ,16GRAM,GPUNVIDIA GeForce GTX3060,显存6G。软件运行环境:Windows1064位操作系统，编程语言为Python3编程，采用基于pytorch深度学习框架。对程序切片需要用到Joern工具，Joern需要用到Linux环境，利用虚拟机实现，虚拟机版本Vmware workstations15.5 Pro，Linux版本Ubuntu18.04.4，内存8G，python版本2.7。切片用到的工具Joern版本0.3.1，Neo4j版本Community-2.1.5，Gremlin 8G，Python-joern 0.3.1,joern-tools 0.3.1,JDK 1.7。

图1是本发明所提方法的流程图，包括代码预处理流程和漏洞检测模型训练测试流程，分别如图2、图3所示。代码预处理流程如图2所示，获取NIST的漏洞数据库，NIST的漏洞数据库包含生产软件漏洞的NVD和包含学术安全漏洞的SARD。在NVD中，每个漏洞都有一个唯一的公共漏洞和标识符(CVE ID)和一个公共的弱点枚举标识符(CWE ID)，该标识符指示着所涉漏洞的类型。随机地选择80％的程序作为训练程序，剩下的20％作为测试程序。从程序中提取C/C++库/API函数调用，有6045个C/C++库/API函数调用，包括标准库函数调用，基本的Windows API和Linux内核API函数调用。总共从程序中提取了56902个库/API函数调用，包括7255个正向函数调用和49647个反向函数调用。将API函数分为5类漏洞问题构建成本文工业软件漏洞的漏洞库，5类漏洞问题分别为缓冲区溢出问题，空指针引用问题，堆溢出问题，API误用和信息泄露问题。本发明采用的5大类漏洞问题的部分关键点如表1所示。

表1部分程序漏洞与包含的关键点

以漏洞库的关键点作为程序切片的切入点提取出代码中与关键点有关的参数、语句和表达式，之后利用工具joern对漏洞库的关键点进行切片，利用数据依赖图，来提取前向切片和后向切片，获得切片后的程序，然后将切片后的程序组装成行关联代码。根据上述5大类问题，将行关联代码数据库总结出5个数据集。5个数据集的漏洞数据类型及行关联代码数量如表2所示。

表2漏洞数据类型及行关联代码数量

对行关联代码进行数据处理，主要是对行关联代码添加标签，包含漏洞的标签设为1，不包含漏洞的标签设为0，对于从NVD程序中提取出的行关联代码，如果行关联代码至少有一行语句通过补丁形式被删除或修改，那么会自动标注为“1”，否则标注为“0”，由于这种自动标注方式可能会误标，所以还需要手动检查自动标注为“1”的行关联代码，并进行修改；对于从SARD程序中提取出的行关联代码，由于每个在SARD中的程序中都已经标注为“good”，“bad”，“mixed”，所以从带有“good”标签的程序中提取出的行关联代码标注为“0”，从带有“bad”或者“mixed”标签的程序中提取出的行关联代码如果至少一行包含导致漏洞的语句，那么该行关联代码标记为“1”，否则标注为“0”。之后进行去重处理，对于行关联代码相同并且标签也相同的代码只保留一个行关联代码，其余的删除，如果行关联代码相同但是标签不同，那么删除这些行关联代码。去重处理完后要删除程序中的注释，并且按行做分词处理，将用户自定义变量名和自定义函数名用“variable_m”、“function_n”进行替换，以免对漏洞语义产生干扰。

然后将行关联代码做分词处理后的单词转换成对应的向量，本发明采用基于层序softmax的word2vec的跳字模型，它基于分布式表示的思想，能将每个词表示成定长的向量，并使这些向量能较好地表达不同词之间的关系。基于层序softmax的方法使用二叉树的数据结构，能大大减少训练的梯度计算开销。跳字模型关心给定中心词生成背景词的概率，利用反向传播更新梯度时，跳字模型用多个词去调整网络，所以在面对生僻词时，跳字模型的效果要优于连续词袋模型。通过训练跳字模型，将损失函数最小化，获得训练样本，利用同样的方法获得测试样本。

由于行关联代码向量化后的向量长度可能不一致，因此，需要将所有的行关联代码向量统一处理为相同长度的向量，当长度不足设置的长度时，则在行关联代码向量最后补0；当行关联代码向量长度超过设置的长度时，超过的部分删除。将行关联代码向量化后，便可将其输入到神经网络中进行训练和检测了。

漏洞检测模型训练测试流程如图3所示，训练过程，行关联代码的向量表示的标志数目tokens设置为50个，dropout设置为0.5，batch size设置为64，epochs设置为300，优化算法采用Adam，选择300个隐藏节点，采用学习率0.001来进行训练。采用批量抽取的方式从训练样本中每次抽取固定个数的样本，传入到采用自注意力机制的代码漏洞检测模型中。

代码漏洞检测模型采用自注意力机制的D_transformer模型，它的网络结构如图4所示，全部采用encoder-decoder架构，由12个堆叠在一起的encoder层和12个堆叠在一起的decoder层组成。每一个encoder和decoder的内部简版结构如图5所示，对于encoder，包含一个自注意力机制层和一个前馈神经网络，自注意力机制能帮助当前节点不仅仅只关注当前的词，从而能获取到上下文的语义。decoder也包含encoder提到的两层网络，但是在这两层中间还有一层encoder-decoder注意力层，帮助当前节点获取到当前需要关注的重点内容。

自注意力机制实际上是想让网络注意到整个输入中不同部分之间的相关性，它的实现方法如下：对于给定的输入向量的每一个向量，分别乘上三个系数矩阵W^q、W^k、W^v获得三个值Q、K、V，分别表示输入query、键key、值value，并且Q＝K＝V。利用点积的方式通过Q和K计算每两个输入向量之间的相似度A，如式(1)所示：

A＝K^T·Q； (1)

对相似度A除以一个常数然后进行加权求和softmax操作后与V利用点积的方式获得自注意力机制层的输出量，输出向量O的公式如式(2)所示：

其中，常数d_k一般指矩阵的第一个维度，起到调节作用，使得点积结果不至于太大。

在自注意力机制中还加入了多头注意力机制，即不仅仅只初始化一组Q、K、V的矩阵，而是初始化多组，本发明采用4组，具体过程为：首先Q、K、V进行线性变换，并做4次缩放点积自注意力，获得4个头，每个头的公式如公式(3)所示：

head_i＝O(QW_i ^q，KW_i ^k，VW_i ^v)； (3)

其中W_i ^q、W_i ^k、W_i ^v分别为Q、K、V第i个头的系数矩阵，W_i ^q∈R^512×128，W_i ^k∈R^512×128，W_i ^v∈R^512×128。

这里要做4次，每次算一个头，而且每次Q、K、V进行线性变换的参数W是不一样的。由于前馈神经网络无法输入4个矩阵，因此需要把4个矩阵连在一起，这样会获得一个大的矩阵，再随机初始化一个矩阵和这个组合好的矩阵相乘，最后获得一个最终多头注意力结果的矩阵作，如公式(4)所示：

MultiHead(Q，k，V)＝Concat(head₁，......head₄)W^o； (4)

其中W^o为随机初始化的矩阵，W^o∈R^512×512。

多头注意力机制进行了4次计算而不仅仅是计算一次，这样可以允许模型在不同的表示子空间里关注不同方面的信息，作为一个集成的作用，提高模型的拟合能力。在自注意力机制中还采用了残差网络中的short-cut结构，目的是解决深度学习中的梯度退化问题。

在把数据输入到前馈神经网络进行激活函数之前，为使输入数据不落在激活函数饱和区，需要进行层标准化LN，LN是在每一个样本上计算均值和方差，公式如式(5)所示：

其中，x_i是输入样本，α是尺度因子、β是平移因子、∈是为了避免除数为0时所使用的微小正数、μ_l是样本数据均值、σ_l ²是样本数据的方差。

D_transformer的详细结构如图6所示，为了让模型捕捉到单词的顺序信息，在encoder和decoder层的输入添加了一个额外的位置编码向量，位置编码的计算方法如公式(6)和(7)所示：

其中，PE为二维矩阵、pos是指当前词在句子中的位置、i是指向量中每个值的index、d_model是词向量的维度。

编码器通过处理输入序列启动。然后将顶部编码器的输出转换为一组注意向量k和v。每个解码器将在其“encoder-decoder attention”层中使用这些注意向量，这有助于解码器将注意力集中在输入序列中的适当位置；完成编码阶段后，开始解码阶段。解码阶段的每个步骤从输出序列输出一个元素。重复此过程，一直到表示解码器已完成输出的符号。每一步的输出在下一个时间步被送入底部解码器，解码器像就像我们对编码器输入所做操作那样，我们将位置编码嵌入并添加到这些解码器输入中，以表示每个单词的位置。解码完成后，在结尾添加一个全连接层和softmax层，获得概率最大的词作为最终结果。

D_transformer突破了RNN模型不能并行计算的限制，并且相比CNN，计算两个位置之间的关联所需的操作次数不随距离增长，因此D_transformer采用完全并行化的结构，训练速度更快，使用较深的模型，较深的模型相比于较浅的模型可以收敛于较低的测试误差，所以该模型可以获得更高的精度。

代码通过漏洞检测模型D_transformer，获得预测结果，将预测结果与真实结果的误差通过反向传播及梯度下降算法进行梯度更新，通过多次迭代训练模型；测试过程，使用训练过的模型进行测试，对比模型测试结果与实际结果是否相同，测试模型的检测能力。

以上所述仅为本发明的优选实施例而已，并不用于限制本发明，对于本领域的技术人员来说，本发明可以有而各种更改和变化。凡在本发明的精神和原则之内，所作的任何修改、等同替换、改进等，均应包含在本发明的保护范围之内。因注意到：相似的标号和字母在下面的附图中表示类似项，因此，一旦某一项在一个附图中被定义，则在随后的附图中不需要对其进行进一步定义和解释。

Claims (8)

1.一种基于自注意力机制的工业软件漏洞检测方法，其特征在于，该方法采用自注意力机制对工业软件漏洞进行检测，检测过程分为代码预处理过程和漏洞检测模型训练测试过程：

所述的代码预处理过程的步骤为：

S11、根据数据库漏洞程序，提取API函数调用和程序片段，将API函数分为5类漏洞问题构建成工业软件漏洞的漏洞库，以漏洞库的关键点作为程序切片的切入点，提取出代码中与关键点有关的参数、语句和表达式；

S12、利用切片工具对漏洞库的关键点进行切片，获得切片后的程序，然后将切片后的程序组装成行关联代码，对行关联代码进行数据处理，并将行关联代码转换成对应的向量，获得特征向量；

S13、将特征向量统一处理为相同长度的向量，当行关联代码向量长度不足时，则在向量后补0，当行关联代码向量长度超过设置的长度时，超过的部分删除；

所述的代码漏洞检测模型训练和测试过程为：

S21、训练过程采用批量抽取的方式从训练样本中每次抽取固定个数的样本，传入到采用自注意力机制的代码漏洞检测模型中，获得预测结果；

S22、将预测结果与真实结果的误差通过反向传播及梯度下降算法进行参数更新，通过多次迭代训练模型；

S23、测试过程使用训练过的模型进行测试，对比模型测试结果与实际结果是否相同，测试模型的检测能力。

2.根据权利要求1所述的基于自注意力机制的工业软件漏洞检测方法，其特征在于，所述数据库采用NIST的漏洞数据库的源代码，包含软件产品中漏洞的NVD和学术安全漏洞的SARD，随机选择80％的程序作为训练程序，20％的作为测试程序。

3.根据权利要求1所述的基于自注意力机制的工业软件漏洞检测方法，其特征在于，所述5类漏洞问题分别为缓冲区溢出问题、空指针引用问题、堆溢出问题、API误用和信息泄露问题。

4.根据权利要求1所述的基于自注意力机制的工业软件漏洞检测方法，其特征在于，所述切片包含前向切片和后向切片，其中前向切片对应于受相关参数影响的语句，而后向切片对应于可影响相关参数的语句，利用数据依赖图，来提取这两种切片。

5.根据权利要求1所述的基于自注意力机制的工业软件漏洞检测方法，其特征在于，所述行关联代码为语义上相互关联的代码行。

6.根据权利要求1所述的基于自注意力机制的工业软件漏洞检测方法，其特征在于，所述数据处理主要是对行关联代码添加标签，包含漏洞的标签设为1，不包含漏洞的标签设为0，对行关联代码做分词处理和去重处理以及针对数据样本不均衡问题为样本赋予不同的权重。

7.根据权利要求1所述的基于自注意力机制的工业软件漏洞检测方法，其特征在于，所述将行关联代码数据转换成对应的向量采用基于层序softmax的word2vec的跳字模型。

8.根据权利要求1所述的基于自注意力机制的工业软件漏洞检测方法，其特征在于，所述代码漏洞检测模型为采用自注意力机制的D_transformer模型，其网络结构采用encoder-decoder架构，内部的encoder和decoder的各包含一个自注意力机制层和一个前馈神经网络。

Images