👌

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

2024/12/22に公開

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

https://github.com/HikaruEgashira/vulnhunters

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

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

vulnhuntrsは、SASTとLLM(大規模言語モデル)を組み合わせることで、これらの課題を解決する新しいアプローチを提案します。

アーキテクチャ

vulnhuntrsは、以下の2つのコアコンポーネントで構成されています:

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

SASTエンジンの実装

SASTエンジンの核となるのは、SymbolExtractorです。以下のような3層構造で脆弱性を検出します:

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. パターンマッチングレイヤー

脆弱性の種類ごとに特化したパターン検出を実装しました:

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との統合

SASTエンジンで抽出された情報は、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")
);

LLMは、この構造化された情報を基に:

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

SASTとLLMの相乗効果

この組み合わせにより、以下のような利点が生まれます:

  1. 高精度な検出

    • SASTによる厳密なパターンマッチング
    • LLMによるコンテキスト理解
  2. 誤検知の削減

    • SASTの構造的な解析
    • LLMによる妥当性確認
  3. 新しい脆弱性への対応

    • パターンベースの検出
    • LLMの推論能力の活用

まとめ

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

Discussion