AI代码审计-skill实践

按我的理解，Agent开发的思路有两种：

• • 带有验证机制的多Agent架构+参数规模比较大的大模型
• • 使用垂直领域的数据集对轻量模型进行微调+简单的Agent架构或是无Agent

前面研究过DeepAudit，我尝试修改了内部的agent架构，比如先做plan（Orchestrator  Agent），然后act（Recon Agent、Analysis Agent、Verification Agent、Reflection  Agent），最后reflect（Reflection  Agent），发现token的消耗变得很多，而且得到的结果其实还是很极大程度取决于大模型的效果，于是我又想到了使用LLAMA  Factory去微调本地的模型，但是又卡在了缺少合适的数据集。

恰巧这两天我发现一直用的iflow cli可以用skill，我希望基于skill这种范式让工具的使用更加的简洁明了，也算是为后续去做AI代码审计提供一些奠基。

目前我想要实现以下skill：

• • 依赖分析skill
• • 源码审计skill
• • 报告输出skill

依赖分析skill

我让iflow实现了依赖分析skill：

目前主要还是基于java，逻辑为优先使用mvn去解析pom文件，如果没有mvn手动则解析pom文件，让iflow调用skill得到总结，依赖版本符合要求，也的确成功调用了mvn：

查看了一下实际输出格式，符合要求：

依赖分析skill主要针对java，但也为其他语言提供了位置以便后续拓展。

源码审计skill

然后我开始让他审计：

实际上都是误报，以下是规则集的问题：

• • source不够全，比如在controller中MultipartFile类型的入参也是可以作为用户输入的，

  

• • sink不够全，因为有可能是直接返回污染而非使用writer去存储污染

  

然后我发现检测逻辑问题更大，仅仅基于sink去做模式匹配：

这样的结果显然不是我想要的，于是重新梳理下思路：

• • skill需要支持多种语言
• • 我需要基于AST树构建调用图，提供source和sink以获取每一对source和sink的完整代码上下文

当收集完所有source-sink对的代码片段后，就可以让LLM基于提示词做污点分析，最后输出报告，这种思路实际上可以不用实现污点追踪。

为了方便，这里我决定直接使用codeql做源码审计

codeql的使用需要两个步骤：

• • 构建数据库
• • 使用ql语句查询

于是我分别创建了codeql create database skill和 codeql query skill

然后在规则集下的java目录中添加了一个文件上传漏洞规则集：

/**
 * @name 文件上传漏洞
 * @description 文件上传漏洞transferTo
 * @kind path-problem
 * @id java/upload
 * @problem.severity error
 * @security-severity 9.2
 * @precision high
 * @tags security
 */

import java
import semmle.code.java.dataflow.TaintTracking
import semmle.code.java.security.PathSanitizer
private import semmle.code.java.dataflow.ExternalFlow
private import semmle.code.java.dataflow.FlowSources
private import semmle.code.java.security.Sanitizers

private class Mapping extends Annotation{
    //出现额外mapping注解在此补充，类似注解只能出现在方法上面
    Mapping(){
        this.getType().getQualifiedName() in
        ["org.springframework.web.bind.annotation.PostMapping",
            "org.springframework.web.bind.annotation.RequestMapping",
            "org.springframework.web.bind.annotation.GetMapping",
            "org.springframework.web.bind.annotation.PutMapping",
            "org.springframework.web.bind.annotation.DeleteMapping",
            "org.springframework.web.bind.annotation.PatchMapping"]
    }
}

private class UploadSource extends DataFlow::Node{
  UploadSource() {
    exists(Method m, Parameter p |
      exists(Mapping mapping | m = mapping.getParent()) and
      p.getCallable() = m and
      this.asParameter() = p
    )
  }
}

private class UploadSink extends DataFlow::Node {
    UploadSink(){
        exists(MethodCall call|(call.getMethod().getDeclaringType().hasQualifiedName("org.springframework.web.multipart", "MultipartFile"))
        and call.getMethod().toString()in["transferTo"]
        and this.asExpr() = call.getArgument(0))
    }

}

module UploadConfig implements DataFlow::ConfigSig {
  predicate isSource(DataFlow::Node source) { source instanceof UploadSource  }

  predicate isSink(DataFlow::Node sink) { sink instanceof UploadSink }

}


module UploadFlow = TaintTracking::Global<UploadConfig>;
// module UploadFlow = DataFlow::Global<UploadConfig>;
import UploadFlow::PathGraph

from UploadFlow::PathNode source, UploadFlow::PathNode sink
where UploadFlow::flowPath(source, sink)
select sink.getNode(), source, sink, "file upload"

报告生成skill

总体效果呈现

现在开始让iflow cli去进行审计：

仅仅使用skill完成对blog-vue-springboot的代码审计，中间过程保存在result目录下

最终结果：

总结

codeql  查询skill  查询套件还有点问题，报告输出也不不够美观，偶尔会出现不使用报告输出skill的情况，但总体能够实现我想要的效果，也不用考虑复杂的多Agent架构，后续也可以对skill进行更多地拓展，比如依赖分析skill如果发现fastjson版本在1.2.83以后，那么实际上codeql  查询skill就没必要去加载fastjson的规则集去查询了