论文复现_Asteria: Deep Learning-based AST-Encoding for Cross-platform Binary Code Similarity Detection

kitsch0x97

已于 2024-07-24 15:24:19 修改

阅读量429

点赞数 3

分类专栏：论文复现文章标签：深度学习人工智能

于 2024-06-03 21:21:50 首次发布

本文链接：https://blog.csdn.net/kitsch0x97/article/details/139422163

版权

论文复现专栏收录该内容

7 篇文章 0 订阅

订阅专栏

在Github上下载论文的代码，并对代码进行分析。这些代码可以实现三个功能，（1）二进制代码的 AST 的生成和预处理；（2）Tree-LSTM 模型的训练；（3）执行漏洞搜索任，代码的目录结构如下图所示。

1. AST 生成功能

进入项目的 ASTExtraction 目录，分析 Asteria 的 AST生成功能。根据帮助文档可知，该功能有两种不同的实现，分别是插件实现和命令行实现，代码的目录结构如下图所示。

首先分析 AST 生成功能的命令行实现，在命令行中输入下面的代码：

# IDA Pro 路径修改为自己的路径
python .\API_ast_generator.py --ida_path C:\Users\22845\Desktop\IDA_Pro_7.7\idat.exe --binary .\vul_database\openssl-1.0.1s --database .\vul_database\vul.sqlite --function EVP_EncodeUpdate

通过 python 中的 pdb 对代码的主要功能进行分析，代码首先会从命令行中解析输入的参数，并将对应的参数打印出来，具体参数如下所示。

Namespace(binary='.\\vul_database\\openssl-1.0.1s', compilation='default', database='.\\vul_database\\vul.sqlite', directory=None, function='EVP_EncodeUpdate', ida_path='C:\\Users\\22845\\Desktop\\IDA_Pro_7.7\\idat.exe', logfile='api_ast_generator.log', timeout=360)

通过输入的参数初始化 AstGenerator 对象，如果 args.directory 为True，则使用 AstGenerator 对象的 extract_ast_from_dir 方法提取对应目录二进制文件的 AST。args.binary 与 args.function 为 True 则使用该对象的 extract_function_ast 方法提取对应二进制文件的AST，具体代码如下所示。

# 创建 AstGenerator 对象，用于提取目标目录或目标文件的 AST
ag = AstGenerator(args)

# 如果 args.directory 为 True 则提取目标目录的 AST
if args.directory:
    ag.extract_ast_from_dir(args.directory)

# 如果 args.binary 与 args.function同时为 True 提取目标文件的 AST
elif args.binary and args.function:
    ag.extract_function_ast(args.binary, func_name=args.function)

根据参数的解析结果binary='.\\vul_database\\openssl-1.0.1s' 与 function='EVP_EncodeUpdate'，可确定程序会使用 AstGenerator 对象的 extract_function_ast 方法生成二进制文件的 AST，其具体实现如下。

def extract_function_ast(self, binary_path, func_name): # the described function 2
    # 定义了一个名为 IDA_ARGS 的列表,其中包含了一些 IDA Pro 命令行参数。
    # 这些参数用于指定编译目标、数据库路径等。
    IDA_ARGS = ["-o %s" % self.args.compilation,  "-g gcc", "-d %s" % self.args.database]
    # 如果用户指定了要分析的函数名,则将其添加到 IDA_ARGS 列表中。
    if self.args.function:
        IDA_ARGS.append("-f %s" % self.args.function)
    # 将 IDA_ARGS 列表中的参数拼接成一个字符串。
    IDA_ARGS = " ".join(IDA_ARGS)
    # 根据操作系统的类型(Linux 或其他),设置适当的环境变量。
    Header = ""
    if platform == 'linux':
        Header = "TVHEADLESS=1"
    # 构建命令行参数列表,包括 IDA Pro 可执行文件路径、命令行参数以及要分析的二进制文件路径。
    cmd_list = [Header, self.args.ida_path, "-c" ,"-A", "-Lida.log" , '-S"%s %s"' % (self.Script, IDA_ARGS), binary_path]
    # 将 cmd_list 列表中的参数拼接成一个字符串。
    cmd = " ".join(cmd_list)
    # 使用 subprocess.Popen 执行该命令
    p = subprocess.Popen(cmd, shell =True)
    try:
        # 设置了一个超时时间,如果命令执行超时,将记录错误日志并终止进程。
        p.wait(timeout=self.args.timeout) # after waiting , kill the subprocess
    except subprocess.TimeoutExpired as e:
        self.logger.error("[Error] time out for %s" % binary_path)
    # 如果命令执行成功(返回码为 0),则记录日志信息;否则,记录错误日志信息并终止进程。
    if p.returncode != 0:
        self.logger.error("[ERROR] cmd: %s" %(cmd))
        if p:
            p.kill()
    else:
        self.logger.info("[OK] cmd %s " % cmd)

执行上述方法，通过命令行调用 ast_generator.py 实现目标二进制文件 AST 的生成，命令行如下所示。

' C:\\Users\\22845\\Desktop\\IDA_Pro_7.7\\idat.exe -c -A -Lida.log -S"C:\\Users\\22845\\Downloads\\Asteria-master\\ASTExtraction\\ast_generator.py -o default -g gcc -d .\\vul_database\\vul.sqlite -f EVP_EncodeUpdate" .\\vul_database\\openssl-1.0.1s'

在上述指令中 “C:\\Users\\22845\\Desktop\\IDA_Pro_7.7\\idat.exe” 是IDA Pro的可执行文件路径，“-c”标志表示在控制台模式运行指令，“-A”标志表示执行自动脚本（批处理模式），“-Lida.log”指定日志文件为ida.log，“-S"C:\\Users\\22845\\Downloads\\Asteria-master\\...."”指定启动后要执行的脚本。这个指令比较长，其单独分析如下：

“C:\\Users\\22845\\Downloads\\Asteria-master\\ASTExtraction\\ast_generator.py”是自定义python脚本的路径。
“-o default”表示脚本的输出格式。
“-g gcc”指定GCC编译器为目标。
“-d .\\vul_database\\vul.sqlite”表示脚本的输出路径。
“-f EVP_EncodeUpdate”指定需要处理的函数的名称

“.\\vul_database\\openssl-1.0.1s”指定被被分析的文件或目录。通过上述指令不难发现，二进制代码 AST 的生成依赖于 “st_generator.py” 文件，该文件的 main 函数如下所示。

if __name__ == '__main__':
    # 处理命令行参数
    import argparse
    ap = argparse.ArgumentParser()
    # 添加编译优化级别相关的命令行参数
    ap.add_argument("-o","--optimization", default="default", help="optimization level when compilation")
    # 添加待处理函数相关的命令行参数
    ap.add_argument("-f","--function", default="", help="extract the specific function info")
    # 添加编译器名称相关的命令行参数
    ap.add_argument("-g","--compiler", default="gcc", help="compiler name adopted during compilation")
    # 添加数据库相关的命令行参数
    ap.add_argument("-d","--database", default="default.sqlite", type=str, help="path to database")
    # 解析传递给脚本的命令行参数，并存储在 args 对象中。
    args = ap.parse_args(idc.ARGV[1:])
    # 创建 AstGenerator 对象，传入参数 optimization 和 compiler。
    astg = AstGenerator(args.optimization, compiler=args.compiler)
    # 调用 astg 对象的 run 方法
    astg.run(astg.get_info_of_func, specical_name=args.function)
    # 创建 DBOP 对象，传入数据库路径参数 database。
    dbop = DBOP(args.database)
    # 将分析结果保存到数据库
    astg.save_to(dbop)
    # 刷新数据库
    del dbop # free to call dbop.__del__() , flush database
    # 退出程序
    Exit(0)

通过 main 函数可以推断 “st_generator.py” 文件的主要功能是，（1）接收命令行参数；（2）根据参数分析二进制文件；（3）将分析结果保存到指定数据库。其中，astg 对象的 run 方法负责分析二进制文件，其代码如下：

def run(self, fn, specical_name = ""):
        '''
        :param fn: a function to handle the functions in binary
        :param specical_name: specific function name while other functions are ignored
        :return:
        '''
        # 如果 specical_name 不为空，则在日志中记录正在处理 specical_name 对应的函数
        if specical_name!="":
            l.info("specific functino name %s" % specical_name)
        # 迭代处理二进制文件中的所有函数，get_func_qty() 用于获取函数的数量
        for i in range(0, get_func_qty()):
            # 获取二进制文件中的第 i 个函数
            func = getn_func(i)
            # 计算并显示当前处理进度
            self.progreeBar(int((i*1.0)/get_func_qty()*100))
            # 获取函数所在段的名称
            segname = get_segm_name(getseg(func.start_ea))
            # 如果段名称的对应位置字符不在指定字符串列表中，则跳过该函数
            # 为什么要执行过滤
            if segname[1:3] not in ["OA", "OM", "te", "_t"]:
                continue
            # 获取函数名称
            func_name = GetFunctionName(func.start_ea)
            # 如果 specical_name 存在且函数名称与 specical_name 不匹配的情况下跳过当前循环
            if len(specical_name) > 0 and specical_name != func_name:
                continue
            try:
                # 分析传入的函数，输出抽象语法树、伪代码、被调用次数、调用次数
                ast_tree, pseudocode, callee_num, caller_num = fn(func)
                # l.error("AST_TREE:"+type(ast_tree))
                # 保存函数名称、起始地址、抽象语法树、伪代码、被调用次数以及调用次数信息
                self.function_info_list.append((func_name, hex(func.start_ea), pickle.dumps(ast_tree), pseudocode, callee_num, caller_num))
            except Exception as e:
                l.warning("%s error" % fn)
                l.warning(str(e))

上述代码的核心功能是从二进制代码中提取目标函数的抽象语法树、伪代码、被调用数量、调用数量等信息，这部分功能由 astg 对象的 get_info_of_func 方法实现，其代码如下：

@staticmethod
def get_info_of_func(func):
    '''
    :param func:
    :return:
    '''
    try:
        # 使用 IDA Pro API 将函数起始地址（func.start_ea）对应的函数进行反编译
        cfunc = idaapi.decompile(func.start_ea)
        # 创建一个 Visitor 对象，并传入反编译后的函数对象 cfunc。
        vis = Visitor(cfunc)
        # 应用 Visitor 对象到函数的主体上，进行遍历或其他处理操作。
        vis.apply_to(cfunc.body, None)
        # 返回抽象语法树、伪代码、被调用数量以及调用数量。
        return vis.root, vis.get_pseudocode(), vis.get_callee(), vis.get_caller()
    except:
        l.warning("Function %s decompilation failed" % (GetFunctionName(func.start_ea)))
        raise

上述代码通过 IDA Pro 的 API 对目标函数进行反编译，并通过 Visitor 对象生成反编译后目标函数的抽象语法树、伪代码、被调用数量以及调用数量。通过执行上述指令，可以生成一个名为 vul.sqlite 的文件，该文件是二进制代码的 AST 表示，如下图所示。