一款方便、高效的基于tree-sitter的代码风格转换器，支持Python、C/C++和Java共100多种风格转换

文章目录

- 工具介绍
- 环境搭建
- 使用教程
- 工具框架流程图
- 增删算法
- tree-sitter介绍
- 项目目录结构说明
- 风格列表

源代码网址：https://github.com/rebibabo/SCTS/tree/main

如果有引用本文或者工具，请注明引用来源

如果觉得对您有帮助，还请各位帅哥美女在github点个免费的star🙏

工具介绍

实现了一款多语言代码等价语义转换器，基于tree-sitter开发工具包，对代码进行解析，并在具体语法树（concrete syntax tree, CST）树上进行节点的匹配和替换，支持Python、C/C++和Java共100多种风格转换，处理2000份代码仅需要13.6s。

代码风格转换器可以应用在以下的场景

后门/对抗攻击：实现隐蔽性高的攻击
数据集扩展：提高模型性能和鲁棒性，使模型更好的对抗扰动
模型水印：向数据集的代码注入水印，保护模型或数据集的产权
统一风格：增强代码的可读性
代码混淆：防止反编译

现有代码风格转换器方法可以大致分为两类：
语法分析器

优点：精确匹配语法规则，保证语义等价和语法正确，转换速度快，能够大批量处理数据集
缺点：需要设计转换规则，比较费时

深度学习模型

优点：不需要指定转换规则，只需要通过数据集训练模型即可，转换成功率高
缺点：token长度限制，生成的代码语法可能错误，并且生成速度非常慢。

本工具收集整理了四种语言C，C++，Python和Java四种语言的112条语义等价变换，设计了增删算法，能够高效处理大批量代码，提升转换速率。

环境搭建

运行命令pip install -r requirements.txt安装python环境，如果想可视化CST树，请安装graphviz，CST树样例如下
在这里插入图片描述
linux请运行sudo apt-get install git graphviz graphviz-doc安装，windows系统请从官网上下载，并配置环境变量
如果不想可视化CST树，可以注释掉所有和graphviz相关的代码

使用教程

构建scts类，设置代码的语言种类

from SCTS.change_program_style import SCTS
scts = SCTS('c')

然后调用change_file_style方法，第一个参数为风格列表，第二个参数为代码文本，返回值第一个元素为风格转换后的代码，第二个元素为转换是否成功，判断依据为代码是否发生变化且语法正确。

new_code, succ = scts.change_file_style(["8.11", "0.1"], code)

例如可以将C语言的for循环修改风格，将初始化语句提到for循环前，条件语句移到循环体内部，自增语句放在循环体末尾。
在这里插入图片描述
其他的风格转换样例

如果想要可视化CST树，执行以下代码

scts.see_tree(code)

如果想将代码进行分词，返回分词列表，执行下面的代码

scts.tokenize(code)

工具框架流程图

工具框架如下图所示，初始化模块输入代码的语言，根据指定语言构建对应的解析器和风格字典的键为风格名称，值为该风格的规则，根据每种编程语言的特点，构建相应的风格字典。风格字典中的键是风格名称，而值则是对应风格的转换规则。例如，对于C语言，风格名称可以是“0.1”，其中小数点左边的数字“0”表示该风格所属的类型，即“修改变量名”，而小数点右边的数字“1”表示该类型的子风格，即“驼峰命名法”。这样就构建了一个风格规则：“0.1: (‘val’, ‘camel’)”，表示将变量名转换为驼峰命名法。风格字典就是由该语言设计的所有风格组成。再输入源代码，经过格式化代码之后，生成CST树，最后输入风格列表，例如[0.1, 1.1, 3.2]，然后查询风格字典，得到所有风格的匹配规则和替换规则。

匹配模块：输入风格列表中所有风格的匹配规则，并初始化匹配节点列表为空，该列表存放所有匹配规则成功的节点。接着输入代码的CST树，递归遍历CST的每一个节点，如果该节点不符合匹配规则，如图中红色实线所示，则遍历该节点的所有子节点，如果符合匹配规则，则将其加入到匹配节点中，根据不同风格的匹配规则，可能继续嵌套的遍历子节点或者不继续往下遍历，直到遍历完所有CST节点，匹配阶段结束，最终输出匹配节点列表。

替换模块：输入风格列表中风格的替换规则以及匹配节点列表，对列表中的每一个节点，根据该风格对应的替换规则，得到修改操作（将在下一节中详细介绍），并将其加入到修改列表中，遍历完成后，根据修改列表，在源代码上直接进行修改，得到该风格转换后的代码。如果还有风格需要修改，则跳回到初始化模块，查询下一个风格的匹配规则和替换规则，重复上述的步骤，直到所有的风格都转换完毕，输出最终的修改后的代码，并检查语法，如果代码发生了改变并且语法正确，则转换成功，否则转换失败。
在这里插入图片描述

增删算法

tree-sitter编辑CST树的成本比较高，而且操作不方便，容易出错，因此考虑对源代码直接进行替换，而不是编辑CST树。在替换模块中，对源代码进行修改操作可以分为插入操作和删除操作两类，通过这两种操作的组合实现替换操作。

为了形式化定义这两种操作，假设原始代码为x，定义插入操作为 I(x, i, s)，表示在字节偏移量为i处的右边插入字符串s，则插入操作如公式3-1所示,而删除操作定义为 D(x, i, j)，表示在字节偏移量为i处，向左删除i个字节：

$I (x, i, s) = x [: i] + s + x [i + 1 :]$
$D\left(x,i,j\right)=x\left[:i-j\right]+x\left[i+1\right]$

所有的修改都可以总结为这两种操作，如果是替换操作，假设想将第i到第j字节偏移的内容修改为字符串s，则替换操作S(x, i, j, s) 可以由删除操作和插入操作组合：
$S\left(x,i,j,s\right)=I\left(D\left(x,j,j-i+1\right),i,s\right)=x\left[:i\right]+s+x\left[j+1:\right]$

所有的修改操作合并组成修改列表O = [O1, O2, … ]，其中Oi表示第i个操作，由于插入或删除会导致接下来的代码内容的字节偏移发生变化，而所有的修改操作都是在原始没有修改的代码上求得的字节偏移量，因此需要设计算法，来弥补每一次修改的字节偏移，并能够准确的对源代码进行修改，设计以下算法，如算法1增删算法所示，输入原始代码和修改列表，输出经过一系列增删操作后修改的代码。

首先需要对操作进行排序，按照修改位置从小到大排序，这样就不会有冲突，另外，由于修改后的偏移发生了变化，因此需要一个变量diff来记录这个偏移量，每一次修改都需要累加diff，例如在当前位置i插入了长度为d的字符串，则diff需要加上d，那么原先在i之后的操作，假设位于j，那么现在的偏移将变成j+diff，即j+d，删除操作diff需要减去删除的字节数，这样就能保证所有替换操作互相之间不会冲突。在算法当中，x[a:b]表示代码x字节偏移a到字节偏移b之间的子字符串，加法操作代表字符串的拼接。

在这里插入图片描述

tree-sitter介绍

在线生成CST树网址：https://tree-sitter.github.io/tree-sitter/playground
在这个网站中，可以解析指定语言，查看各个节点的属性和对应的代码，便于可视化解析
在这里插入图片描述

该库会对指定的语言生成解析器，并保存在build/下的.so文件中，后使用Language类加载该.so文件，得到对应语言的编译器，接着生成具体语法树CST，具体请参考change_program_style.py文件下的SCTS.__init__(self, langauge)方法。
CST节点的固有属性如下表所示
在这里插入图片描述
通过节点固有属性，我们可以获取到节点和源代码之间的相关性，例如行号、代码等，然而对于节点的操作不仅限于固有属性，还需要能够遍历CST树，获取节点与节点之间的关系，下表为CST节点的关系属性，通过获取节点的关系属性，可以按照不同顺序遍历。
在这里插入图片描述

项目目录结构说明

dataset目录下存放测试用的代码
c、cpp、java、python目录下存放各个语言的风格转换器，transform*.py文件中包含rec_*函数，用于匹配节点，cvt_*函数用于替换节点，config.py为配置文件

add_transform.py文件可以增加自定义的风格，需要自己写代码（参考transform*.py)，执行python add_transform.py --l java --n type可以增加java语言type类风格的转换器

confuse.py可以混淆python代码，执行效果
原始代码如下

def can_add_ball(x, y): # 当前坐标是否覆盖了其他桌上的球
    for i in coord:
        if dot2dot(x, y, i[0], i[1]) < 2 * ball:
            return False
    return True

def balls(): # 添加球
    for _ in range(int(np.random.randint(8,15))):
        x, y = np.random.uniform(ball, length - ball), np.random.uniform(ball, width - ball)
        if can_add_ball(x, y):
            ax.add_patch(mpatches.Circle([x, y], ball, color='red'))
            coord.append([x, y])
    for color in ['blue', 'yellow', 'green', 'brown', 'black', 'pink']:
        x, y = np.random.uniform(ball, length - ball), np.random.uniform(ball, width - ball)
        if can_add_ball(x, y):
            ax.add_patch(mpatches.Circle([x, y], ball, color=color))
            coord.append([x, y])

混淆代码如下

def f_36(t_11, t_22): 
    for t_28 in t_42:
        if (f_38.__call__(t_11, t_22, t_28[((((14 % 5 + 3) * 2 - 10) // 4) - 1)], t_28[((5 * 4 - (6 + 2) + 9) // 20)]) < (((9 * (4 - 2) + 3) % 10)*2) * ball) and (sum(x for x in range((5*(((20 // 4) + 6 - (3 ** 2)) // 2)))) > (((4 * 3 - (5 + 1) + 8) // 10) - 1)):
            return not (all(type(x) is int for x in [((8 * (3 - 1) + 5) % 4), (((3 ** 3 // 9 - 5 + 4) // 2)*2), (((2 ** 3 + 10 // 2 - 9) % 3)*3), ((((20 // 4) + 6 - (3 ** 2)) // 2)*2*2)]))
    return (True if (((20 // 4) + 6 - (3 ** 2)) // 2)+(((3 * 4 - 6) // 2 + 5) // 8) == (((5 * 4 - (6 + 2) + 9) // 20)*2) else False) and (False if ((((15 % 4) ** 2 - 3) // 5)*2)*(2*(((5 * 3 - 10) // 2 + 1) // 3)) != (2*2*((10 * (3 - 1) + 5) % 4)) else True)

def f_10(): 
    for t_31 in range.__call__(int.__call__(np.random.randint.__call__((2*((2 ** 3 - 4 * 2 + 8) // 8)*2*2),(5*((10 * (3 - 1) + 5) % 4)*3)))):
        t_11, t_22 = np.random.uniform.__call__(ball, length - ball), np.random.uniform.__call__(ball, width - ball)
        if (f_36.__call__(t_11, t_22)) and (bool(dict(a=((2 ** 3 + 10 // 2 - 9) % 3), b=((((5 * 3 - 10) // 2 + 1) // 3)*2)))):
            ax.add_patch.__call__(mpatches.Circle.__call__([t_11, t_22], ball, t_12='\x72\x65\x64'))
            t_42.append.__call__([t_11, t_22])
    for t_12 in ['\x62\x6c\x75\x65', '\x79\x65\x6c\x6c\x6f\x77', '\x67\x72\x65\x65\x6e', '\x62\x72\x6f\x77\x6e', '\x62\x6c\x61\x63\x6b', '\x70\x69\x6e\x6b']:
        t_11, t_22 = np.random.uniform.__call__(ball, length - ball), np.random.uniform.__call__(ball, width - ball)
        if (f_36.__call__(t_11, t_22)) and (sum(x for x in range((((10 * (3 - 1) + 5) % 4)*5))) > (((10 * (3 - 1) + 5) % 4) - 1)):
            ax.add_patch.__call__(mpatches.Circle.__call__([t_11, t_22], ball, t_12=t_12))
            t_42.append.__call__([t_11, t_22])