CN110008710A - 基于深度强化学习和程序路径插桩的漏洞检测方法 - Google Patents

Abstract

一种结合深度强化学习和程序路径插桩技术的漏洞检测方法，首先通过插桩的方式从待测程序的控制流图中得到输入所对应的路径，根据路径及控制流图中的目标节点计算得到奖励值，然后将奖励值用于训练深度强化学习的神经网络，从而用于选择变异动作，根据该变异动作对待测程序的输入进行变异后得到更新后的输入及其路径并计算更新后的奖励值并再次训练神经网络和进行输入变异处理，循环至待测程序崩溃，即得到相应的输入漏洞。本发明准确率更高，并可以更高效地获得漏洞所在路径对应输入，相较于传统的模糊测试，检测速度更快，兼具一定的代码覆盖量。

Description

基于深度强化学习和程序路径插桩的漏洞检测方法

技术领域

本发明涉及的是一种信息安全领域的技术，具体是一种基于深度强化学习(DQN)和程序路径插桩的漏洞检测方法。

背景技术

漏洞是在硬件、软件、协议的具体实现或系统安全策略上存在的缺陷，使得攻击者不经允许访问和修改计算机。现有漏洞分析技术主要有人工检测、Fuzz技术、补丁比对、静态分析和动态分析。

发明内容

本发明针对现有漏洞分析技术的局限和不足，提出一种基于深度强化学习和程序路径插桩的漏洞检测方法，以自然语言处理(NLP)方法将样本抽象出的向量作为特征，以程序执行路径到目标代码块距离作为奖励，训练深度学习网络，以此网络作为样本变异器。本发明准确率更高，并可以更高效地获得漏洞所在路径对应输入，相较于传统的模糊测试，检测速度更快，兼具一定的代码覆盖量。

本发明是通过以下技术方案实现的：

本发明首先通过插桩的方式从待测程序的控制流图中得到输入所对应的路径，根据路径及控制流图中的目标节点计算得到奖励值(reward)，然后将奖励值用于训练深度强化学习的神经网络，从而用于选择变异动作，根据该变异动作对待测程序的输入进行变异后得到更新后的输入及其路径并计算更新后的奖励值并再次训练神经网络和进行输入变异处理，循环至待测程序崩溃，即得到相应的输入漏洞。

所述的待测程序是指：可编译执行的完整代码而非代码片段。

所述的控制流图是指：待测程序的抽象表现，它代表了一个程序执行过程中会遍历到的所有路径，该控制流图优选通过Soot工具得到。Soot是McGill大学的Sable研究小组自1996年开始开发的Java字节码分析工具，它提供了多种字节码分析和变换功能，通过它可以进行过程内和过程间的分析优化，以及程序流图的生成，还能通过图形化的方式输出，让用户对程序有个直观的了解。

所述的插桩的方式是指：在保证逻辑完整性的基础上在待测程序中插入若干探针，通过探针采集运行过程中产生的特征数据，从而获得待测程序的控制流信息和数据流信息并进一步得到逻辑覆盖动态信息，从而实现测试目的的方法。

所述的插桩，优选通过工具PIN实现。Pin是Intel推出的一款动态二进制插桩工具，目前在学术界使用十分广泛。一般用于程序动态分析用，并且支持多个平台，windows、linux以及OSX。

所述的特征数据是指：打印语句输出的结果，即通过探针在每一代码块的入口处或出口处添加计数语句或打印语句，对语句执行状态进行结果输出。

所述的控制流信息是指：描述程序逻辑执行的先后顺序的信息。

所述的数据流信息是指：描述程序运行过程中数据的流转方式及其行为状态的信息。

所述的逻辑覆盖动态信息是指：测试过程中所得到的覆盖率信息，通过对程序逻辑结构的遍历实现程序的覆盖，它是一系列测试过程的总称。

所述的目标节点是指：控制流图中已经确定或疑似的漏洞代码和代码块和/或敏感函数所在的结点，该目标节点通过预先标记方式得到，例如该节点所对应的代码块包含敏感函数，则将该节点标记为目标节点。

所述的奖励值是指：根据当前输入对应的路径与目标节点，根据计算出的路径到目标节点的距离，奖励值与距离反相关，其中：d_s(s,T_b)为种子文件的执行路径与目标节点的距离；s为种子文件seed；T_b为目标节点；ξ_b(s)为种子文件s作为输入时，待测程序的完整执行路径；m为执行路径中的某一节点，d_b(m,T_b)为节点m与目标节点的距离。

所述的深度强化学习的神经网络采用Deep Q-Learning Network(DQN)，该神经网络的结构内设有长短期记忆(LSTM)模型层与四层全连接层，该神经网络将输入经过word2vec处理后得到的向量作为LSTM模型层的输入，LSTM模型层的输出则作为四层全连接层的输入，全连接层输出为当前输入下各动作对应Q值，即Q(S,A)。

所述的神经网络的初始参数经随机处理，通过训练过程中对于每个输入-动作-奖励值-新输入(S,A,R,S’)元组，使用Q-learning算法进行网络参数的更新。

所述的神经网络中采用经验回放与固定Q目标的方式建立与评估网络结构相同的目标网络，实时更新评估网络的参数而间歇更新目标网络的参数，以降低样本之间的相关性。经验回放指从以往的经验即经验库中随机采样进行训练。经验回放的步骤为：在经验库经验达到200条后，5次获得新经验，从经验库内随机抽取32条经验，使用Q-learning算法更新网络参数。固定Q目标指每获得200次经验后，将Q网络中的评估网络的参数完全赋给更新网络。

所述的变异动作包括：随机bit位翻转(random_bit_flip)、byte级洗牌(shuffle)、插入令牌(insert_token)、增减随机数(arith)、随机byte/word/dword替换(interest)。

所述的崩溃是指：待测程序经理非正常退出、无响应、报错、输出无法控制以及意料之外的输出，造成崩溃的输入称为崩溃输入。

本发明涉及一种实现上述方法的漏洞检测系统，包括：预处理模块、奖励模块、内置Deep Q-Learning(DQN)模型的训练模块和输入变化模块，其中：预处理模块以二进制格式读入原始种子，并运用自然语言处理的方法将其抽象成一组向量作为特征，奖励模块根据待测程序的控制流图、程序执行路径与目标节点计算得到奖励值，训练模块根据奖励值对神经网络进行训练并得到选择变异的动作，输入变化模块根据动作对原始种子进行变异处理得到新种子并输出至预处理模块进行迭代的漏洞检测。

所述的种子是指：预制且在实际测试中导致待测程序崩溃的文件。

所述的自然语言处理的方法包括但不限于卷积算法等。

技术效果

本发明基于NLP、DQN算法和程序插桩技术，以自然语言处理(NLP)方法将样本抽象出的向量作为特征，以程序执行路径到目标代码块距离作为奖励，设计并实现变异方法，对程序进行模糊测试，效果明显。相对于传统的模糊测试工具，具有更高的导向性和效率，同时具有一定的代码覆盖率。对于漏洞检测准确度可以达到76％。

附图说明

图1为本发明流程示意图；

图2为LSTM-Q网络结构图；

图3为LSTM基础结构。

具体实施方式

本实施例针对的漏洞代码可为跨平台语言，主要适用语言为C/C++。如图1所示，为本实施例涉及的一种基于深度强化学习(DQN)和程序路径插桩的漏洞检测系统，包括：预处理模块、变异模块、奖励模块和训练模块，其中：预处理模块根据模糊测试的种子生成一组特征，特征经过神经网络获得变异动作，变异模块根据动作变异种子，生成新种子，奖励模块结合待测程序的控制流图，程序执行路径与目标节点进行计算得到奖励值，训练模块根据奖励模块反馈的奖励值对Deep Q-Learning(DQN)模型进行训练。

所述的漏洞代码库为从Sard数据库中搜集得到的缓冲区溢出代码。

本实施例以任意文件或实现准备好的种子文件作为模糊测试种子，此处预先准备好的种子文件可以是在实际测试中总结出的大概率导致待测程序崩溃的文件。例如以一个.pdf文件作为种子文件，以二进制方式打开，读取其01bit串作为待测程序的输入。

步骤一：系统构建，具体包括：

1.1)训练NLP网络，将模糊测试种子映射为实数域上的向量，并作为LSTM网络输入；

所述的NLP网络包含一个输入层、一个隐藏层、一个输出层，输入层输入为种子文件，输出为32维向量，表示的是其32维特征，并作为LSTM网络的输入。NLP网络的参数已经预先训练过。

1.2)搭建LSTM-Q评估网络以及结构相同LSTM-Q目标网络并随机初始化两个网络的参数，具体为：建立一个包括一层输入层、一层隐层、一层输出层的LSTM网络，该LSTM网络的输出作为Q网络的输出；另外建立一个具有两个神经网络的Q网络，分别称为评估网络和目标网络，两个网络结构相同，由四层全连接层构成，四层的神经元数目分别为32,64,180,8，其中LSTM网络与评估网络的组合称为LSTM-Q评估网络，LSTM网络与目标网络的组合称为LSTM-Q目标网络，LSTM和Q网络的参数均为随机初始化。

1.3)构建经验库，具体为：将大小为200的数组用于储存上文的(S-A-R-S’)元组或称经验，每次获得经验后，将存储入经验库中，当经验库已满即经验数目达到200后，再次获得经验将替换掉经验库中存储时间最长的经验。

1.4)使用PIN对待测程序插桩，使用Soot生成待测程序控制流图并根据敏感函数所在程序块确定目标节点。获取待测程序模糊测试种子。种子为预先准备的种子文件。

所述的敏感函数是memcpy.strcpy等可能造成缓冲区溢出的系统函数。

步骤二：变异种子，具体包括：

2.2)将种子输入NLP网络得到实数域变量，以实数域变量作为特征分别输入LSTM-Q网络得到该特征下各动作Q值，即状态-动作值Q(s,a)。

2.2)根据epsilon-greedy策略从变异动作中选取随机动作或最优动作a，具体为：以epsilon的概率选取随机动作，以1-epsilon的概率选取最优动作，最优动作是指最大的Q(s,a)值所对应的动作a。

所述的随机动作选取概率epsilon优选为0.1。此处epsilon也可采用阶段降低策略，具体为初始epsilon为0.5，每获得200次经验后降低0.1，直至最低0.1。

步骤三：储存经验，具体包括：

3.1)以最优动作a对原种子进行变异操作得到新种子并输入待测程序，计算程序执行路径到目标节点距离并得到报酬r。

3.2)将新种子输入NLP网络得到对应的新状态s’，构建四元组(S,A,R,S’)作为一条经验储存在经验库内。

所述的经验库大小优选为200。

所述的储存优选为队列方式，即经验库满后获得经验将替换先前储存经验。

步骤四：更新网络，具体包括：循环执行步骤二至步骤三，并当经验库内经验数目达到200后，每5次获得经验，采用分批学习方式，从经验库中随机抽取32条经验，并根据Q-learing算法更新评估网络；在获得200次经验后，将Q网络中的评估网络的参数完全赋予目标网络。

所述的Q-learing算法是指：Q(s,a)←Q(s,a)+α[r+γmaxQ(s′,a)-Q(s,a)]，其中：报酬衰减因子γ为0.9，学习因子α为0.5，s,a,r,a’分别指代上文中四元组(S,A,R,S’)元素。

所述的更新评估网络是指：将loss函数值传递给评估网络，并更新评估网络参数，其中Loss函数L＝E[(TargetQ(s,a)-Q(s,a))²]，TargetQ(s,a)为目标网络得到的Q(s,a)值。

优选地，当步骤二至步骤四运行一段时间后，将使待测程序崩溃的种子记录到文档内，便于人工分析待测程序是否存在漏洞。

本方法通过使用强化学习DQN的方法，使种子的变异过程具备导向性，避免种子的无效或重复变异，如始终对某byte进行random_bit_flip变异动作，变异前后程序执行路径不变，从而提高模糊测试的平均速度。

使用Sard数据集与同种模糊测试工具AFL进行对比测试结果如下，本发明具备在基本保持代码覆盖率的情况下，一定程度上提升模糊测试速度的特点：

相同执行时间下代码覆盖率：为AFL的97.2％，相比AFL略低。

同一种子文件、待测程序下下执行时间：为AFL的89.4％，速度提升10.6％。

上述具体实施可由本领域技术人员在不背离本发明原理和宗旨的前提下以不同的方式对其进行局部调整，本发明的保护范围以权利要求书为准且不由上述具体实施所限，在其范围内的各个实现方案均受本发明之约束。

Claims (7)

1.一种基于深度强化学习和程序路径插桩的漏洞检测方法，其特征在于，首先通过插桩的方式从待测程序的控制流图中得到输入所对应的路径，根据路径及控制流图中的目标节点计算得到奖励值，然后将奖励值用于训练深度强化学习的神经网络，从而用于选择变异动作，根据该变异动作对待测程序的输入进行变异后得到更新后的输入及其路径并计算更新后的奖励值并再次训练神经网络和进行输入变异处理，循环至待测程序崩溃，即得到相应的输入漏洞；

所述的插桩的方式是指：在保证逻辑完整性的基础上在待测程序中插入若干探针，通过探针采集运行过程中产生的特征数据，从而获得待测程序的控制流信息和数据流信息并进一步得到逻辑覆盖动态信息，从而实现测试目的的方法；

所述的目标节点是指：控制流图中已经确定或疑似的漏洞代码和代码块和/或敏感函数所在的结点，该目标节点通过预先标记方式得到，例如该节点所对应的代码块包含敏感函数，则将该节点标记为目标节点；

所述的奖励值是指：根据当前输入对应的路径与目标节点，根据计算出的路径到目标节点的距离，奖励值与距离反相关，其中：d_s(s，T_b)为种子文件的执行路径与目标节点的距离；s为种子文件seed；T_b为目标节点；ξ_b(s)为种子文件s作为输入时，待测程序的完整执行路径；m为执行路径中的某一节点，d_b(m，T_b)为节点m与目标节点的距离。

2.根据权利要求1所述的方法，其特征是，所述的深度强化学习的神经网络采用DQN，该神经网络的结构内设有LSTM模型层与四层全连接层，该神经网络将输入经过word2vec处理后得到的向量作为LSTM模型层的输入，LSTM模型层的输出则作为四层全连接层的输入，全连接层输出为当前输入下各动作对应Q值，即Q(S，A)。

3.根据权利要求2所述的方法，其特征是，所述的神经网络的初始参数经随机处理，通过训练过程中对于每个输入-动作-奖励值-新输入(S，A，R，S′)元组，使用Q-learning算法进行网络参数的更新。

4.根据权利要求2所述的方法，其特征是，所述的神经网络中采用经验回放与固定Q目标的方式建立与评估网络结构相同的目标网络，实时更新评估网络的参数而间歇更新目标网络的参数，以降低样本之间的相关性，其中：经验回放指从以往的经验即经验库中随机采样进行训练。

5.根据权利要求4所述的方法，其特征是，所述的经验回放具体为：在经验库经验达到200条后，5次获得新经验，从经验库内随机抽取32条经验，使用Q-learning算法更新网络参数，固定Q目标指每获得200次经验后，将Q网络中的评估网络的参数完全赋给更新网络。

6.根据权利要求1所述的方法，其特征是，所述的变异动作包括：随机bit位翻转、byte级洗牌、插入令牌、增减随机数、随机byte/word/dword替换。

7.一种实现上述任一权利要求所述方法的漏洞检测系统，其特征在于，包括：预处理模块、奖励模块、内置DQN模型的训练模块和输入变化模块，其中：预处理模块以二进制格式读入原始种子，并运用自然语言处理的方法将其抽象成一组向量作为特征，奖励模块根据待测程序的控制流图、程序执行路径与目标节点计算得到奖励值，训练模块根据奖励值对神经网络进行训练并得到选择变异的动作，输入变化模块根据动作对原始种子进行变异处理得到新种子并输出至预处理模块进行迭代的漏洞检测；

所述的种子是指：预制且在实际测试中导致待测程序崩溃的文件。