写了个 LangChain 驱动的自动代码审计工具

写在前面

最近 DeepSeek 真是火的一塌糊涂，这也让我不得不思考，未来是否真的会被 AI 取代？

我想，与 AI 紧密配合也许会是未来的大方向：某些机械性的工作不会再需要那么多的人力去完成，更多时候应该是交由大模型驱动的工作流来自动完成，而仅需少量人员作为监督和审查。

随着安全意识的不断提高，SDL 也越来越被重视，我就在想：SDL 中的 SAST 是否也可以让 AI 来辅助完成？在网上搜索了一番，确实有部分团队已经在做这个事情了，但是数量并不多，所以我决定尝试一下，把这一想法付诸实践。但此前我对大模型一无所知，所以整个过程也都磕磕绊绊的，最终写出来的工具也仅限于能用，还有很多不足，希望未来能深入学习之后再做提高吧。

项目地址： https://github.com/yv1ing/MollyAudit

前置知识

对于大模型本身而言，它是无法与外部的数据进行交互的，早期的 ChatGPT 也仅仅能根据用户的输入来预测要输出的内容。大模型的输出完全由训练数据来决定。但后来随着大模型技术的不断迭代，出现了能将外部数据引入大模型中，使得大模型不仅能回答通用问题，还能根据外部数据，针对性地解决某些特定领域的问题。

LangChain 就是为了实现这一目标而出现的由语言模型驱动的应用程序框架，它的核心组件有：

• Compents： 为 LLMs 提供接口封装、模板提示和信息检索索引；
• Chains： 将不同的组件组合起来构成一个链条，按照链条上的次序完成特定的任务；
• Agents： 使得 LLMs 能够与外部环境进行交互，例如通过 API 请求执行操作；

简单来说，LangChain 实现了两点：

1. 数据连接： LangChain 允许 LLMs 连接到外部数据源，并与之进行交互；
2. 行动执行： LangChain 允许 LLMs 根据特定指令来完成特定动作；

工作流程

LangChain 的工作流程如下图：

1. 首先将本地的文档数据通过非结构加载的方式转成文本，然后再通过分词器转换成文本块，再通过嵌入器将文本块转换成向量进行存储（RAG）;
2. 用户输入信息，通过嵌入器转换成向量，在向量数据库中检索，并通过提示词的形式输入 LLMs，最终通过 LLMs 得到答案；

文本嵌入

文本嵌入（text embedding）是一种将高维数据映射到低维空间的技术。在处理文本时，每个单词或短语都被转换为一个固定大小的向量。这些向量可以捕捉到原始文本中的语义和语法信息，使得具有相似意义的单词在向量空间中彼此接近。

向量化之后的文本存储在向量数据库中，当需要输出信息到 LLMs 时，LangChain 会先在向量数据库中检索相关的信息，一并输入 LLMs，使得 LLMs 的输出更具相关性。

工具设计

简单理解了 LangChain 的工作机制之后，就开始了工具的设计，毕竟深入学习的时间成本比较高，我暂时也不需要，所以够用就好。

我主要使用 LangChain 提供的数据连接能力，把需要审计的项目提取出项目结构和源代码集两部分，通过 system prompt 输入 LLMs 具体的动作规范，然后将项目结构输入 LLMs，由 LLMs 自行决定审计哪些源代码，审计的过程信息向量化之后存储在向量数据库中，作为后续审计的参考数据。为了避免上下文超出 LLMs 的限制，顺便使用了一下 LangChain 提供的上下文压缩链，参考上下文压缩

实际编码虽然磕磕绊绊，但也只花了不到两天，毕竟功能也不算复杂，目前实现了：

1. 图形化界面（pyqt6）
2. 自定义 Base API 和 API key；
3. 导入项目目录；
4. 导出审计结果；

为了使 LLMs 能够按照预期完成自动审计的任务，我在 system prompt 中约定了 4 个动作：

1、查询项目结构：

<root>  
    <action>QUERY STRUCTURE</action>  
    <content></content>  
</root>

2、查询指定文件：

<root>
    <action>QUERY SOURCE</action>
    <content>[absolute path]</content>
</root>

3、输出审计结果：

<root>
    <action>OUTPUT RESULT</action>
    <content>[result content]</content>
</root>

4、结束审计任务：

<root>
    <action>FINISH TASK</action>
    <content></content>
</root>

目前测试来看，这四个动作已经能较好地完成中小型项目的审计任务。

未来的一些设想：

1. 收集整理代码审计相关的文章，作为知识库输入 LLMs，提高审计的准确率；
2. 收集整理已知漏洞的 Poc 和 Bypass 样本，由 LLMs 反推出漏洞成因，构建一个样本库，同样作为后续审计的参考；
3. ……

使用截图

在开发过程中，我选用了几个开源的代码审计靶场来进行测试，审计过程全自动完成，检出效果也还不错，属于能用的类型。