前言

上个月编写Semgrep规则时的想法是先利用Semgrep规则进行初步代码审计，然后结合GPT来辅助确认代码审计结果，打算实现一个结合代码扫描和生成式人工智能的自动代码审计系统。
早期有很多项目利用GPT进行代码扫描，但由于Token的限制，实际应用中并不是那么便利。一般只能识别一些简单漏洞，同时消耗的Token也相当可观，并且效果还没那么好。如果先用静态扫描确定sink，那么再用gpt进行辅助代码审计，这样效果就会好很多。

具体实现

要实现这个想法，主要就两个点，一个是用来识别漏洞代码片段的静态扫描工具，第二个就是用于辅助代码审计的大语言模型。

格式的选择

首先第一点，对于静态扫描工具，当然是希望能兼容的工具越多越好，所以在了解了静态扫描工具的输出后，SARIF格式输出成为了一个最好的选择，目前市面上流行的静态扫描工具都支持SARIF格式的输出。
静态分析结果交换格式(SARIF) 是一种开放标准，旨在简化不同工具和系统之间静态分析结果的交换，用于在不同工具之间交换和理解静态分析报告。它最初是为了满足Eclipse插件开发者的需求而设计的，现在也被广泛应用于其他静态代码分析工具中。SARIF其实就是json格式，目前SARIF 2.1.0已经是OASIS标准。

静态扫描工具的选择

对于静态扫描工具的选择，我选择个人最常用的Semgrep和Codeql来进行分析。
Semgrep的规则编写简单易学，虽然可能会出现误报，但结合LLM可以显著减少误报的概率。通过结合LLM，可以有效地改善Semgrep的准确性，使其输出更可靠。
在另一方面，CodeQL是一个强大的工具，其规则执行后能够提供源代码，为LLM的代码审计提供辅助，特别是在污点追踪方面。通过结合CodeQL进行污点追踪，可以使LLM的审计结果更加精确，有助于深入了解漏洞和安全问题的来源。

llm的选择

llm我直接选择了常用的gpt3.5 32k的api，目前大多数大模型都支持100k+的token识别，但是考虑到价格的因素，最终还是选择了gpt3.5的api。当然自己去微调训练一个代码审计模型效果理论上来说只要数据质量够高，效果也不会差到哪去，甚至可能更高。

整体设计框架

系统实现

由于在使用Codeql时，看到了l3yx师傅的Choccy项目，UI非常符合我的审美，并且也十分好用，于是打算直接魔改这个项目，改下前后端代码逻辑，加入Semgrep扫描和调用llm进行分析和展示的代码。
用以前写的go靶场测试，最后实现的效果就是这样
Semgrep扫描和LLM分析展示页面

Codeql扫描和LLM分析展示页面

prompt的艺术

在没有进行微调的情况下，构建 prompt 变得尤为关键。对于我个人在构造 prompt 时，考虑到 Semgrep 没有污点分析功能，因此我会直接将 sink 点周围的代码加入，或者如果 token 允许的话，我会将包含 sink 的整个代码文件添加到 prompt 中以进行分析。
而对于 CodeQL，由于其支持污点分析，我会将 sink 附近的代码和相关 source 代码一同加入到 prompt 中。这有助于模型更好地理解数据流和可能存在的漏洞路径，提高审计的深度和全面性。
除了代码之外，我还会将扫描文件的规则描述添加到 prompt 中，以辅助生成式模型的分析。这一步能够提供额外的上下文信息，帮助模型更好地理解审计背景和需求，进而提高审计结果的准确性和可靠性。
通过整合各种信息和数据来构建 prompt，可以提升生成式模型对代码审计场景的理解和处理能力，为自动代码审计系统的开发和应用带来更显著的效果和价值。
然而这样就够了吗？这里给出一些优化方案。

SAST的规则描述的优化

对于 SAST 规则，我们可以进一步拓展其应用。举例来说，对于修复了 XXE 的 Java 代码，若考虑到阻止外部实体的需求，Semgrep 编写规则可能会相对复杂。不过，LLM 可以轻松解决这一问题，只需在 Semgrep 规则中添加补充说明 XXE 的修复方法，以增强 LLM 的识别能力。
同样地，针对企业内常用的修复方法，可利用 prompt 来增强 LLM 的识别能力。每家企业都有独特的代码框架和一套修复方法，因此编写适合自身的规则就是 prompt 的妙处。
通过这样的方法，我们可以使 LLM 在审计过程中更具针对性，帮助识别特定修复方案中的漏洞和安全隐患。个性化的 prompt 不仅提升了模型的识别能力，还促进了自动化审计系统与企业内部开发实践的契合。

xmlReader.setFeature("http://apache.org/xml/features/disallow-doctype-decl", true);
xmlReader.setFeature("http://xml.org/sax/features/external-general-entities", false);
xmlReader.setFeature("http://xml.org/sax/features/external-parameter-entities", false);

factory.setProperty(XMLConstants.ACCESS_EXTERNAL_DTD, "");
factory.setProperty(XMLConstants.ACCESS_EXTERNAL_SCHEMA, "");

factory.setProperty(XMLInputFactory.SUPPORT_DTD, false);
factory.setProperty("javax.xml.stream.isSupportingExternalEntities", false);

输入代码的优化

除了输入source和sink，其实还可以考虑一些其他的。打个比方，针对fastjson漏洞，一般我们在白盒扫描中为了提高准确率，编写的其中一个规则是这样的：是否引入了fastjson这个包，autotype的配置，是否配置黑白名单，对于白名单的危险类检查等。这样我们可以扩展，比如引入的包等可以辅助llm分析的代码，都可以作为prompt用作漏洞的识别，再结合SAST规则描述的进一步优化，强化prompt的表达能力。

SAST规则的负优化?

有了 LLM 的辅助，我认为并非简单追求 SAST 规则更加准确就是最佳选择。举一个例子，以针对 CodeQL 的 Go SSRF 规则为例，测试中可能会发现有些 sink 判断存在漏报问题，某些不太常见的写法可能被疏忽。如果没有 LLM 的支持，我可能会优化 CodeQL 规则。但在有了 LLM 的辅助后，我可以考虑让 CodeQL 规则更偏向宽泛一些（特别是针对 sink 的判断）。
这种方法的好处在于能够保留 source 的污点分析功能，同时在 sink 方面更容易发现更多漏报场景。借助 LLM 的辅助，我们可以确保降低误报的可能性。通过权衡规则的精确性与宽泛性，来提高审计的覆盖范围并减少漏报情况。
因此在这种情况下，静态扫描工具去兼容不同框架来实现污点分析可能比实现一个高精确的 sink 发现规则是更为合适的选择。

只是静态扫描？

在我的看法中，LLM 的优势在于其能够处理一些通过编写代码比较棘手的任务，也就是理解能力。因此，LLM 的辅助判断在安全领域有许多应用。举个例子，对于增量 API 的提取，有时编写规则来识别代码中不规范编写的 API 是相当繁琐的。在这种情况下，LLM 可能是一个很好的辅助工具。
除此之外，在安全领域中，LLM 的应用还有许多方面。在我看来，只要涉及到需要理解能力的地方，LLM 都是一个很好的选择。特别是通过agent来实现一些流程化的工作，可以是未来的一个方向