通过结合静态代码分析和大语言模型(LLM)的方式来检测代码中的潜在漏洞,据作者描述通过该工具在AI bug bounty平台挖掘出多个漏洞,核心是来学习该项目的结果,如何组织LLM和SAST,同时对该项目的缺点进行分析便于进行定制化的重构
(图源protect ai)
本着阅读项目比阅读文章更能学到新东西的态度对项目进行了分析
首先我们来看下作者提供的一个完整的框架图
大致表达了其完整的workflow
:
下面从代码层面学习其实现方式:
其项目结构如下:
vulnhuntr/
├── .devcontainer/ # 开发容器配置
├── vulnhuntr/ # 主要源代码目录
│ ├── __init__.py
│ ├── __main__.py # 主入口文件
│ ├── LLMs.py # LLM模型相关实现
│ ├── prompts.py # 提示词模板
│ └── symbol_finder.py # 代码符号提取器
├── .env.example # 环境变量示例文件
├── Dockerfile # Docker构建文件
├── pyproject.toml # 项目配置文件
└── requirements.txt # 依赖包列表
从代码分析来看,这是一个使用 LLM (Large Language Model) 进行代码安全分析的工具,主要功能包括:
在其入口函数,执行的是run()
方法
首先看一下框架可选的参数列表
在指定了待分析的项目根目录后,vulnhuntr
将会对项目抽象成RepoOps
类对象,以及SymbolExtractor
类对象
RepoOps
类其中封装了项目操作相关的多个API,主要是在整个工作流中扮演者预处理的角色,详细实现可见后文子模块的分析
而对于SymbolExtractor
类,其同样封装了多个API操作,核心是利用的jedi
这一个python语言下的静态分析框架进行项目代码分析,主要特点为采用多层次搜索策略,确保能够寻找到不同情况的符号定义,同时提供了比较灵活的文件搜索和过滤机制,具体的分析见后面的子模块分析
接下来进入具体的分析逻辑
他会通过封装的RepoOps
类对象的get_relevant_py_files
方法来获取,不包括以下内容的所有项目中的.py
文件
接下来进入一个小小的分支,将会对通过--analyze
参数指定的文件路径或者文件进行漏洞分析,若没有对其进行指定则将会分析完整的项目
不管是使用的get_files_to_analyze
或者get_network_related_files
,其均是返回待分析的文件列表,其中后者是筛选出以下source点文件进行分析
接下来则为通过LLM的能力进行README文件信息的提取,构建出一个System Prompt
出来使得LLM能够更好的理解待分析的项目
其中的README_SUMMARY_PROMPT_TEMPLATE
提示词是用来进行README文件的信息提取
而SYS_PROMPT_TEMPLATE
提示词是用来指定LLM的任务,列举了需要分析的漏洞列表
上述过程的pipeline中得到了如下的关键信息
后续根据上述得到的信息进行进一步的漏洞检测任务
这里待分析的文件包含有两个维度,其一为通过--analyze
参数指定的文件或者文件夹下的文件,另外则是所有network-related
文件(像是Flask框架的@app.route(xx)
类似的API调用文件)
核心的漏洞检测任务将会遍历获取的所有文件列表中的任意个.py
首先当然是加载理解了项目信息的LLM
之后进行初步的分析任务:
这里使用<file_code>
标签包裹待分析的单个.py
文件的所有代码
使用<instructions>
标签包裹我们的进行代码分析的指令
指令剖析:指定了需要检查的sinks,同时表明如果存在安全限制的情况下需要对其进行bypass操作,在这个基础上从API endpoints
为起点进行漏洞检测,同时表明了对于不能够明显的判断其存在对应的漏洞的情况下,可以通过<context_code>
包裹你需要请求的类或者方法代码进行进一步的漏洞分析验证
使用<analysis_approach>
标签用来包裹在漏洞分析过程中的具体方法
指令剖析:
source
点位置流出的数据流信息,需要分析每一类数据流的处理方式、存储位置、输出内容,防止在漏洞检测过程中出现过高的假阳性context_code
域进行上下文函数调用的相关代码记录,在分析过程中进行结合分析只用<guidelines>
包裹一些输出的指南
最后就是通过<response_format>
包裹我们预期的输出格式
核心内容为pydantic定义的类Response
会将上述的所有标签及标签中的内容喂给LLM进行理解分析进行初步的漏洞审查
**总结:**初步审查的目的是用来进行可能的漏洞点存在的筛选,这里不存在有具体的漏洞检测任务
将会对初始审查结果中的可能存在漏洞风险的代码进行对应漏洞类型的具体审查,具体来说分成了多个轮次
对于第一个轮次,其不会进行具体的漏洞审查工作
**个人理解:**这样设置的作用个人感觉主要是出于以下原因进行考虑
initial analysis
进行分析,再根据需要请求更多的上下文好的,下面看一下其如何实现:
在第一个轮次中,并不会请求context code
,同时也不会引入initial analysis
的分析结果
构建具体漏洞审查的prompt
引入待分析.py
文件代码
第一个轮次中不存在context code
,其他轮次中,将会加载获取的context functions and classes
加载一些特定漏洞的Bypass示例,更好的知道LLM进行Bypass任务
引入特定漏洞的具体如何审查的指令
这里以RCE
漏洞为例:
最后规定了LLM生成回复范式,同样包含有下面这几类元素
其他轮次相比于第一个轮次多了两点:1. context code
的加载 2. previous analysis
的协同分析
context code
,在确保了没有请求已经存在的context code
的情况下,将会调用SymbolExtractor
类的extract
方法进行对应Function或者Class的提取,具体如何实现的后续分析context code
之后,同第一轮次一样构建vuln_specific_user_prompt
,在引入了需要的context信息以及上次分析结果的情况下进行特定种类漏洞的审查任务整个项目对于结束分析的规定包含有三类情况
在进行for i in range(7)
的轮次循环的过程中,如果没有新的上下文代码的请求时,也即是获取到了所有的上下文代码信息,则将此时的分析的结果作为最终的分析结果进行输出
同时请求相同的上下文代码两次
最后当然是达到了规定的最大迭代次数7
该类核心是用来进行项目代码相关操作的
get_readme_content
方法:其通过关键词匹配的方式获取项目描述,将其喂给LLM,使得大模型能够充分的了解项目,进行更具针对性的分析
get_relevant_py_files
方法:通过预定义无关文件名以及文件夹的方式进行待分析文件的初步筛选
get_network_related_files
方法:通过预定义的各种Web框架的API路由标识,来获取存在外部可访问的文件信息
该类主要作用为
主要是使用了jedi
这一个Python语言的静态分析框架
https://github.com/davidhalter/jedi
其实现的检索核心是在extract
方法中,其实现了检索范围从小到大的三层搜索策略
文件级的搜索
按照搜索语句的不同类型进行特定种类的检索
项目级的搜索
若上述两种方式均为得到搜索结果,则采用全局名称的方式进行上下文代码的搜索
通过拉取github项目进行了尝试
可以学习的部分:
多层次的分析策略
使用静态分析框架进行上下文代码的检索使得填充上下文代码
prompt提示词的编写
个人感觉的不足之处:
https://protectai.com/threat-research/vulnhuntr-first-0-day-vulnerabilities
https://github.com/protectai/vulnhuntr
https://github.com/davidhalter/jedi
https://jedi.readthedocs.io/en/latest/docs/api-classes.html#abstract-base-class
3 篇文章
如果觉得我的文章对您有用,请随意打赏。你的支持将鼓励我继续创作!