Zenn
👌

SAST + LLMによる次世代の脆弱性診断

2024/12/22に公開
Security
脆弱性診断
sast
LLM
tech

vulnhuntrsというツールを作成しました。GitHubに公開しています。

https://github.com/HikaruEgashira/vulnhunters

静的アプリケーションセキュリティテスト（SAST）は、コードの脆弱性を発見するための重要なツールです。
しかし、従来のSASTツールには以下のような課題がありました：

  • 誤検知が多い
  • コンテキストを理解できない
  • 新しい脆弱性パターンへの対応が遅い

vulnhuntrsでは、SASTとLLM（大規模言語モデル）を組み合わせてこの課題に向き合ってみました。
PoCとなる元リポジトリはこちらです。

アーキテクチャ

vulnhuntrsは、以下の2つのコンポネントで構成されます：

  1. Tree-sitterベースのSASTエンジン
  2. LLMによるコンテキスト分析エンジン

SASTエンジンの実装

SASTエンジンはSymbolExtractorにて定義されています。以下のような3層構造でLLMに与えたい定義部分を検出します：

impl SymbolExtractor {
    pub fn extract(&mut self, name: &str, code_line: &str, files: &[PathBuf]) -> Option<CodeDefinition> {
        for file_path in files {
            if let Ok(definitions) = self.parser.parse_file(file_path) {
                // 1. 関数定義の検索
                if let Some(def) = self.find_function_definition(...) {
                    return Some(def);
                }
                
                // 2. 脆弱性パターンの検索
                if let Some(def) = self.find_vulnerability_pattern(...) {
                    return Some(def);
                }
                
                // 3. 一般的なパターンの検索
                if let Some(def) = self.find_general_pattern(...) {
                    return Some(def);
                }
            }
        }
        None
    }
}

1. 構文解析

Tree-sitterを利用することで複数の言語で対応可能な拡張性の高い設計を目指しました、現在以下の言語に対応しています：

  • Python
  • JavaScript/TypeScript
  • Java
  • Rust
  • Go

2. パターンマッチング

脆弱性の種類ごとに危険な関数(Sink)をパターンマッチングで検出します：

fn find_vulnerability_pattern(&self, definitions: &[Definition], pattern_name: &str, ...) {
    let relevant_patterns: Vec<&str> = if pattern_type.contains("sql") {
        vec!["sql.call", "sql.exec", "sql.method"]
    } else if pattern_type.contains("command") {
        vec![
            "cmd.call", "cmd.exec", "cmd.method",
            "vuln.object - subprocess",
            "vuln.method - system",
            // ...
        ]
    } else if pattern_type.contains("xss") {
        vec![
            "dom.method - innerHTML",
            "vuln.method - render",
            "template.render",
            // ...
        ]
    }
    // ...
}

3. コンテキスト抽出

関数呼び出しチェーンや変数の使用状況を追跡しています：

fn find_function_definition(&self, definitions: &[Definition], name: &str, ...) {
    let name_variations = vec![
        clean_name.clone(),
        format!("function.name - {}", clean_name),
        format!("method.name - {}", clean_name),
        // コンテキストに応じた名前のバリエーション
    ];
    // ...
}

LLMとの統合

最終的に与えるプロンプトはこのようになります。

let prompt = format!(
    "File: {}\n\nContent:\n{}\n\nContext Code:\n{}\n\nVulnerability Type: {:?}\n\nBypasses to Consider:\n{}",
    file_path.display(),
    content,
    context_code,
    vuln_type,
    vuln_info.bypasses.join("\n")
);

以下のような評価軸で出力します：

  1. コードの意図を理解
  2. 潜在的なバイパス方法を考慮
  3. 脆弱性の深刻度を評価

なぜSASTが重要なのか

ソースコードをそのまま与えてもいいと感じるかもしれません。SASTをコンテキスト抽出に利用しなのには以下の理由があります：

  1. 高精度な検出

    • 厳密なパターンマッチングが可能
    • 脆弱性は直感に反した部分に現れるので本当に必要な部分のみに着目することで見つける可能性を上げる

  2. 誤検知の削減

    • SASTの検出はかなり高精度です。人間及びLLMが見逃すような検出もプログラムなら漏れなく検出できます。
    • その分誤検知が多いというトレードオフに対してLLMによる妥当性確認を行うことでお互いの弱点を補完できます

  3. 新しい脆弱性への対応

    • SASTはライブラリ固有の実装には対応できるが、未知のライブラリにはその実装を見ない限り対応できません
    • LLMがパターン認識することで未知のライブラリに対応したSASTを実現できます

まとめ

LLMは4oレベルですでに人間以上の能力を持っていると思います。あとはいかに必要な情報をコンテキストとして与えらるかかなと思ってるので得意領域のSASTを組み合わせたcontext生成を行っているvulnhuntrのリポジトリを参考にしてみました。reflectionなどagenticな処理をふんだんに使ってます。ぜひ触ってみてください。

Discussion