人工智能与漏洞识别-devign【2】训练自定义数据集

训练NIST数据集

https://samate.nist.gov/SARD/testsuite.php

使用Juliet Test Suite for C/C++测试套件

以C为例

对数据集进行slice(切片)

1. 观察数据集结构，以及pdf说明，bad函数名以文件名+"_bad"组成

• 

2. good函数以多种方式命名，需要设定多个规则

• 
  

识别方法：以static void good来进行识别或者void good等多种规则

生成json格式数据集

原来的数据集json含有四个key,

project项目名称，默认为FFmpeg,如果修改，则需要在main.py作相应修改
commit_id Hash值（不重要）
target 1代表有漏洞，0为无漏洞
func具体函数

• 在main.py执行过程中，project和commit_id不会被读入，该值对程序无任何影响,可以随机生成

实现

• 该脚本自带函数改名，并全部改为func,如需要保留源函数名，可以删除相关语句
• 作者无法确定函数名对于神经网络的影响，在测试过程中发现，保留函数名或改为func都会造成最后训练结果Accuracy和F1都为1的奇怪现象，如果采用去掉函数名的方法，accuracy降为0.75，似乎正常了
• 使用该脚本前需要新建dataset.txt在运行目录，提前输入[，并把相应c文件放在同一目录下，运行结束后把最后的,删除

import json
import os
strart =0#状态控制
n_1=0#左花括号计数
n_2=0#右花括号计数
word=""
#打开dataset
data = open("dataset.txt", "a")
#打开c文件
for dirpath, dirs, files in os.walk('.'):                # 递归遍历当前目录和所有子目录的文件和目录
    for name in files:                                   # files保存的是所有的文件名
        if os.path.splitext(name)[1] == '.c':
            filename = os.path.join(dirpath, name)
            word=""
            with open(filename, 'r') as raw:
                list1 = raw.readlines()
                #好的
                for i in list1:
                    if "good" in i :
                        stt=i
                        list2 = stt.split(" ")
                        if list2[0] in ["static", "void"]:
                            print()
                            for i in range(len(list2)):
                                if "good" in list2[i]:
                                    indexx = list2[i].find("(")
                                    list2[i] = "func" + list2[i][indexx:]
                            stt = ' '.join(list2)
                            word+=stt
                            strart=1
                        continue
                    if "goodG2B" not in i and strart==1:
                        word = word + i
                        if "{" in i:
                            n_1 += 1
                        if "}" in i:
                            n_2 += 1
                        if n_2 == n_1 and n_2 != 0:
                            strart = 0
                            n_1 = 0
                            n_2 = 0



                if word!="":
                    word = {'project': 'FFmpeg', 'commit_id': '973b1a6b9070e2bf17d17568cbaf4043ce931f51', 'target': 0,
                            'func': word}
                    json_s = json.dumps(word)
                    data.write(json_s+",")
                    print("good")
                word=""#释放word

                #坏的,target1
                name=name[:-2]
                for i in list1:
                    if "void "+name+"_bad" in i :
                        stt = i
                        list2 = stt.split(" ")
                        if list2[0] in ["static", "void"]:
                            for i in range(len(list2)):
                                if "bad" in list2[i]:
                                    indexx = list2[i].find("(")
                                    list2[i] = "func" + list2[i][indexx:]
                            stt = ' '.join(list2)
                            word += stt
                            strart = 1
                        continue
                    if strart==1 and  "void "+name+"_bad" not in i :
                        word = word+i
                        if "{" in i:
                            n_1+=1
                        if"}"in i:
                            n_2+=1
                        if n_2==n_1 and n_2!=0:
                            strart=0
                            n_1=0
                            n_2=0

                if word != "":
                    word = {'project': 'FFmpeg', 'commit_id': '973b1a6b9070e2bf17d17568cbaf4043ce931f51', 'target': 1,
                            'func': word}
                    json_s = json.dumps(word)
                    data.write(json_s+",")
                    print("bad")
                word=""


#添加结束符
data.write("]")
data.close()