用于代码缺陷检测的数据集生成

论文地址：Towards security defect prediction with AI

代码仓库地址：sa-bAbI

数据集sa-bAbi的生成

核心文件：

sa_babi/
.
├── generate.py  # 主要生成代码的文件
├── sa_tag.py  # 每行代码性质标签
└── templates.py  # 生成不同模块(控制流、函数体等)的模版

如果要添加漏洞种类，则则需要往sa_tag.py中添加新的行标签，往templates.py中添加新的模块模版。

generate.py重要函数逻辑如下：

main()函数逻辑

1. 获得参数和路径并检查
2. 设置种子
3. 得到生成器数组generators（含有各种例子生成器的数组）
   
   • 如果设置为linear_only表示代码顺序执行，不包含控制流，生成器只有一种。
4. 开始循环生成代码例子
   
   1. 利用取模操作获得当前使用的生成器
   2. 传递参数调用对应的生成器生成代码和对应标签
   3. 利用哈希生成文件名，发生碰撞则重新生成代码
   4. 写文件元数据到元数据数组
   5. 将生成代码和标签写入c文件
5. 将元数据数组（该数据集的所有元数据）写入json文件。

各生成器函数逻辑

有以下几种生成器，对应相应的函数：

buffe write 格式："$buf_var[$idx_var] = '$char';"

fv含义：free variable

以上生成器逻辑几乎相同，只是需要的变量和判断safe的条件不一样。逻辑如下：

1. 得到所有匿名变量——调用_get_anon_vars()生成得到。
2. 选前3个变量作为关键变量，其余为无关紧要的变量dummy_vars，只用来传递数值，实际不起什么作用。
3. 通过random随机生成关键变量需要用到的关键数值。
4. 得到buffer write的字符——调用_get_char()生成得到。
5. 用substitutions字典将其封装好。
6. 从模版中得到模块主要代码行templates.xxx_MAIN_LINES。
7. 用关键数值判定是否安全，布尔变量safe用来指示。
8. 从模版中得到模块的声明初始化变量代码行templates.xxx_DEC_INIT_PAIRS。
9. 传入以上参数，调用_assemble_general_example组装模块生成代码。

gen_tautonly_linear_example函数由于不经过控制流，直接进行buffer write，因此所有变量都是dummy_vars，substitutions、main_lines、dec_init_pairs均为空。

如果要添加漏洞种类，我觉得可以按照以上函数的模式写gen_xxx_{Vulnerability Name}_example7个对应函数，具体是否要写那么多个要看漏洞的性质，main函数可能只需要少量改动。

_assemble_general_example()

组装模块函数逻辑：

1. 判定是否包含条件buffer write，包含则给main_lines加一句templates.BUFWRITE_LINES。
2. 调用_get_lines得到所有代码行。（这个函数及相应调用函数需要修改。）
3. 调用_get_tags得到所有代码行标签。
4. 调用_get_instance_str得到最终c语言代码文本。

_get_lines()

获得所有代码行，包含声明初始化、主体内容、dummy变量

（只用来传递数值，实际不起什么作用）的使用。函数逻辑：

1. 调用_get_setup_lines获得声明初始化代码行。
2. setup_lines + main_lines构成所有代码行。
3. 为所有代码行打上body标签，如果包含条件buffer write(include_cond_bufwrite == True)，则为最后一个代码行（即为buffer write代码）打上Tag.BUFWRITE_COND_SAFE或 Tag.BUFWRITE_COND_UNSAFE（通过safe判断）。
4. 在指定最大最小值范围随机生成dummy变量的个数。
5. 调用_insert_dummies为现有代码行添加dummy变量相关的代码行。

_insert_dummies()

插入所有dummy变量的声明和赋值语句。函数逻辑如下：

1. setup_lines + main_lines构成所有代码行。
2. 获得控制流代码行的开始和结束索引。
3. 循环num_dummies次，调用_insert_referential_dummy依次插入每一个dummy变量的相关代码行。

_insert_referential_dummy()

插入指定dummy变量的声明、初始化和buffer write语句。这里面出现的buffer write不经过控制流，因此标签为BUFWRITE_TAUT_SAFE 或 BUFWRITE_TAUT_UNSAFE。

函数逻辑如下：

1. 判断dummy变量个数，小于2报错（一次需要用到两个dummy变量）。
2. 根据是否需要安全赋值，设置dummy数组长度dum_len和buffer write需要用到的下标dum_idx。
3. 从dummy_vars中取出两个变量用作：buffer write的数组变量——dum_buf_var；指示buffer write位置的下标——dum_int_var。
4. 按照格式利用上述的4个参数设置声明、赋值、buffer write代码行。
5. 组织声明初始化代码行顺序。
6. 决定是否声明初始化代码行开始和结束位置。
7. 在开始和结束范围内随机生成声明初始化代码行位置索引。
8. 在声明初始化代码行后的位置随机生成buffer write代码行位置。
9. 根据上述索引修正更新控制流开始和结束位置。
10. 将上述代码行插入原来的代码行lines中。
11. 根据dum_idx、 dum_len的大小关系判断buffer write是否安全。
12. 将上述代码行对应标签插入原来的标签body_tags中。

_test()

测试每个生成器生成代码。

这个函数大部分不需要修改，只需要改从745-757行测试标签种类的代码。

不需要修改的函数

补充-生成*.c.tok文件

*.c.tok文件是按json格式组织的，在tokenizer工具中生成的，是在执行完生成源文件数据后才进行的步骤。

记忆神经网络代码

首先明确深度学习神经网络里面的一些概念：

• epochs：训练趟数；
• batch_size：批大小；
• steps_per_epoch：每一趟的步数，或者理解为每趟的批数。

以上三者的关系是：一般情况下，一趟会训练一次整个的训练集；在一趟中分steps_per_epoch步进行，一步处理batch_size个数据。即 $batch\_size \times steps\_per\_epoch = num\_samples$ 训练集数据总数。

• train set、valid/validation set、test set 训练集、验证集、测试集；
  
  • 训练集用来估计模型，验证集用来确定网络结构或者控制模型复杂程度的参数，而测试集则检验最终选择最优的模型的性能如何。验证集是模型训练过程中单独留出的样本集，它可以用于调整模型的超参数和用于对模型的能力进行初步评估，在每趟结束后会进行评估，不参与训练。
• 独热编码(one-hot coding)：又称一位有效编码，其方法是使用N位状态寄存器来对N个状态进行编码，每个状态都由他独立的寄存器位，并且在任意时候，其中只有一位有效。（适合特征并不总是连续值，而是分类值，并且没有次序关系）通俗的来说，比如，我们有3个类：苹果、雪梨、香蕉，只有一种特征，那么我们就用一个1维长度为3的数组表示这三个类，[ 1., 0., 0.]表示苹果，[ 0., 1., 0.]表示雪梨，[ 0., 0., 1.]表示香蕉。这样便于代码使用。用np.argmax()就可获得最大值的下标，由于数组只有0、1，取最大值就取到了1所在的索引，其实就获得了标签索引。

train.py 训练网络

逻辑如下：

1. 做好数据准备工作：
   
   1. 设置好工作目录
   2. 调用utils.load_data()从npy文件中获得代码实例矩阵instances_mat、标签矩阵labels_mat和 训练集与验证集字典索引partition。
   3. 获得代码实例中的最大行数。
   4. 调用datagen.DataGenerator()得到数据生成器data_generator。
   5. 从数据生成器中获得分类总数目num_classes。
2. 调用run_experiments开始训练。

run_experiments函数

逻辑如下：

1. 获得训练参数：每趟的训练步数、验证步数。
2. 调用datagen.DataGenerator()得到验证集数据生成器val_data_generator。
3. 先调用datagen.generate_balanced()从验证集数据生成器中得到平均了每个类分布(每个类的分布是相同的)后的生成器，再调用next返回该生成器的下一个项目，即获得验证集代码实例矩阵val_instances_mat、验证集标签矩阵val_labels_mat。
4. 调用np.argmax()可获得验证集标签索引数组y_true。
5. 调用np.zeros()初始化混淆矩阵(误差矩阵)cnf_matrices。
6. 按照规定的实验次数循环开始训练网络：
   
   1. 调用juliet_memnet.get_model()获得未训练的网络model。
   2. 编译网络。
   3. 训练网络，30趟。
   4. 用网络预测验证集得到预测结果predics
   5. 调用sklearn.metrics.confusion_matrix()获得预测结果和标签y_true的混淆矩阵并将其存入数组中。
   6. 保存模型到指定文件夹。

datagen.py 获得数据生成器

juliet_memnet.py 构建记忆神经网络

预处理工作

==注意：只要import juliet_memnet 就会执行这段代码！！==

1. 调用utils.generate_sa_data()生成sa_babi数据（就是这一步会调用utils.save_data()生成npy文件）
2. 调用utils.load_data()从npy文件加载获得代码实例矩阵instances_mat、标签矩阵labels_mat和 训练集与验证集分割后的字典索引partition。
3. 调用datagen.DataGenerator()得到数据生成器data_generator。
4. 还有一些参数的获取。

utils.py 处理数据

预处理数据：

1. from sa_babi.sa_tag import Tag获得所有标签。
2. 为后续要得到的CSV文件定义简单token名称。
3. 定义后续要写入的各种文件名称，如npy文件和pkl文件。
4. 定义一些特殊的token。

generate_sa_data函数

生成sa_babi数据矩阵并且保存到工作目录中。

输入参数：

• working_dir ：保存数据的工作目录，默认是sa-train-1000
• coarse_labels：一般不管他。布尔类型，如果是true，表明只有安全和不安全两种情况，默认为false。

逻辑如下：

1. 调用get_examples()从tokens/*.c.tok获得实例instance、标签labels、路径paths。
2. 判断是否是只有TAUT的情况，如是，要做重新映射。(感觉如果加入了我们写的多种漏洞，这里可能不需要了，或者要修改)
   
   1. 调用set()获得不重复的标签集合unique_elements。
   2. 获得只有TAUT情况的标签集合tautonly_elements。
   3. 通过以上两个集合是否相等来赋值给布尔变量is_tautonly。
   4. 如果是coarse_labels或者is_tautonly，重新映射labels里的标签，即不区分TAUT和COND，只保留safe和unsafe。
3. 重新修改标签，body标签也改为other，即只保留了与漏洞有关的标签。
4. 调用get_vocab_mapping()获得词汇映射为整数的字典。
5. 调用get_data_dimensions()获得数据的维数：实例数目、最大行数、最大行长。
6. 调用get_example_matrices()获得实例矩阵、标签矩阵。
7. 调用get_partition()获得训练集与验证集分割后的字典索引partition。
8. 调用print_data_stats()打印数据。
9. 调用save_data()将实例矩阵、标签矩阵、整数映射字典、分割数据集的字典、路径存入npy文件和pkl文件中。

get_examples函数

从tokens/*.c.tok获得实例instance、标签labels、路径paths。

返回值含义如下：

1. instance[i][j][k]是第i个文件第j行第k个token字符串。

   instances: list of list of list of str
   instances[i][j][k] represents
       - the i-th CWE121 C example
       - the j-th line
       - the k-th token in that line

2. labels[i][j]是第i个文件第j行对应标签的整数值。

labels: list of list of int
    labels[i][j] represents
        - the i-th CWE121 C example
        - the j-th line
    in SA: the value is the integer Tag value

1. paths这些*.c.tok*文件的路径数组。

该函数逻辑如下：

1. 获得*.c.tok*文件的路径数组path_list。
2. 循环遍历每个文件路径：
   
   1. 调用ann_tok_to_lines()得到代码行数组lines。
   2. 遍历代码行数组lines，获得包含代码行号和tokens的实例代码数组instances：
      
      1. 排除空的行
      2. 添加行号token
   3. 获得标签数组labels：
      
      1. 由于*.c.tok*与对应的 *.c文件名称相同，因此通过*.c.tok*文件名加上路径拼接找到对应的 *.c文件路径。
      2. 传入 *.c文件路径为参数，调用get_sa_tags()获得sa_babi数据集标签数组tags。
      3. 由于*.c.tok*文件中不包含include行，因此移除相应的标签。
      4. 将标签数组tags中标签对应整数构成最后所需的标签。

get_example_matrices函数

将instances、labels转换为以0填充的np数组。

返回值含义如下：

1. instances_mat：代码实例矩阵，3维np数组，各维度最大值：[num_examples, max_numlines, max_linelen]
2. labels_mat ：标签矩阵：2维np数组，各维度最大值： [num_examples, max_numlines]。 labels_mat[i][j] 是第i个文件第j行的标签

函数逻辑如下：

1. 调用get_data_dimensions()获得数据的维数：实例数目、最大行数、最大行长。
2. 调用get_vocab_mapping()获得词汇映射为整数的字典。
3. 初始化instances_mat、labels_mat。
4. 三层循环遍历instances中的token，来给instances_mat中token对应位置赋映射字典中对应的整数值；同时在两层循环的地方，将labels对应的值直接复制给labels_mat对应的值，注意此时的labels已经合并了body和other标签。

get_partition函数

获得训练集与验证集分割后的字典索引partition。

逻辑如下：

1. 打乱实例索引顺序，idx_list为打乱后的实例索引数组。
2. 确定训练集中样本数目
3. 定义字典partition

字典结果如下：

partition = {
        'train': idx_list[:num_train],
        'validation': idx_list[num_train:]
}

ann_tok_to_lines函数

将*.c.tok文件转换为代码行数组lines。

*.c.tok文件格式如下：

{
    "filename": "/mnt/data/src/0a8b6849a9.c",
    "tokens": [
        {
            "kind": "Keyword",
            "line": 2,
            "sem": "FunctionDecl",
            "text": "int"
        },
        {
            "kind": "Identifier",
            "line": 2,
            "sem": "FunctionDecl",
            "sym": {
                "id": "c:@F@main",
                "kind": "FunctionDecl",
                "type": "int ()"
            },
            "text": "main"
        },
        ....
    ]
}     

函数逻辑如下：

1. 打开文件，加载json格式，只获得['tokens']部分数据tok_data。
2. 获得行数。
3. 初始化代码行字典line_data，该字典长这样——{行号：token字符串数组}。注意这里行号从1开始，但是为了统一从0开始索引，规定line_data[0]=[]。
4. 循环遍历tok_data，将token加入字典中对应行号的数组里。
5. 按行号顺序用字典的值——即token字符串数组构建lines。

get_vocab_mapping函数

获得从词汇元素(代码里的token)到整数的字典映射。

逻辑如下：

1. 获得实例文件里的所有词汇(token)vocab，有重复。
2. 调用sorted(list(set(vocab)))获得无重复并且排序后的词汇。
3. 按顺序建立字典映射，注意保留0不用，索引从1开始。

save_data函数

将实例矩阵、标签矩阵、整数映射字典、分割数据集的字典、路径存入npy文件和pkl文件中。

1. 实例矩阵、标签矩阵存入npy文件（np数组）。
2. 整数映射字典、分割数据集的字典、路径存入pkl文件（两个字典、一个数组）。

load_data函数

将实例矩阵、标签矩阵、整数映射字典、分割数据集的字典、路径从npy文件和pkl文件中加载出来。

不太重要的中间函数

1. print_data_stats()：打印数据。
2. get_data_dimensions()：获得数据的维数：实例数目、最大行数、最大行长。
3. 调用get_sa_tags()获得sa_babi数据集标签数组tags。

未用到的函数

validate.py 验证/预测模型

运行 python validate.py 可能出现的问题和解决方法

1. 找不到matplotlib模块——因为在enviroment.yml里被原作者注释了…，所以在sa_babi环境下pip intall matplotlib。

2. 在virtualenv环境下使用matplotlib绘图时遇到了这样的问题：

>>> import matplotlib.pyplot as plt
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
...
in <module>
from matplotlib.backends import _macosx
RuntimeError: Python is not installed as a framework. The Mac OS X backend will not be able to function correctly if Python is not installed as a framework. See the Python documentation for more information on installing Python as a framework on Mac OS X. Please either reinstall Python as a framework, or try one of the other backends. If you are Working with Matplotlib in a virtual enviroment see 'Working with Matplotlib in Virtual environments' in the Matplotlib FAQ

解决方案来自https://www.cnblogs.com/harelion/p/5637767.html：

似乎是因为虚拟环境与默认环境的安装配置不同造成的。

搜索错误信息之后，在STO上找到了解决方案：

1、pip安装matplotlib之后，会在根目录下产生一个.matplotlib的目录：cd ~/.matplotlib

total 280
-rw-r--r-- 1 me staff 78K 10 4 2015 fontList.cache
-rw-r--r-- 1 me staff 59K 1 17 15:56 fontList.py3k.cache
drwxr-xr-x 2 me staff 68B 10 4 2015 tex.cache

2、在这个目录下创建一个名为matplotlibrc的文件，内容是：

backend: TkAgg

然后保存退出，重启Python交互界面或重新运行脚本，import正常执行。

STO答案地址：http://stackoverflow.com/questions/21784641/installation-issue-with-matplotlib-python

1. 似乎跑到后面还会报错，还没解决....