世界初の脆弱性AI探索をgoで実装し直した話

## TLDR;

1. [vulnhuntr](https://github.com/protectai/vulnhuntr/tree/main?tab=readme-ov-file) という AIによる脆弱性探索を見つけた
2. python しか対応していなかったから、go に対応するように実装し直した（[成果物 vulnhuntr-go](https://github.com/YmBIgo/vulnhuntr-go)）
3. JavaScript, TypeScript, Java や Ruby での実装も挑戦したい

## LLM 群雄割拠の時代。Next Big Things を探せ！

皆さん。LLM 使いこなせていますか？
きっとこの記事に辿り着けるだけの情報ネットワークを持っている皆様であれば、LLM 使いこなせてるよ！って人も多いと思いますが、折角なら使う側ではなくて作る側に回りたいと考えたことはないでしょうか？

でも、Cursor や Github Copilot や Devin などアメリカの天才たちが血眼になって開発をしている中で今からそこに食い込むのは少し難しいかもしれないですね。

ただ考え直したら、今から１年ちょっと前の LLM を使ったコーディングの世界はこんなにも盛況だったでしょうか？
自分の感覚だと今ほどのものではなかったです。なんなら去年の5月頃の自分は OpenDevin（現 [OpenHands](https://github.com/All-Hands-AI/OpenHands) ）を手元で動かしてみて、「AIこんなもんか」とたかを括っていました。

時代が移り変わるのに、１年もかからないなんて恐ろしい業界です。IT業界は。

---

そんな出遅れが怖いけど、だからと言って何から始めれば今からでもキャッチアップできるか分からないあなたに「次のネタ」を提供できればいいなと思って記事を書きました。

結論から述べましょう。個人的に次に来そうな技術的なネタは、

***AIによる脆弱性自動探索***

です！

なぜか？それは、

1. AIによる脆弱性自動探索のPOCはまだ***去年の10月***に発表されたばっかり([vulnhunhr](https://github.com/protectai/vulnhuntr/tree/main?tab=readme-ov-file))
2. POCなのでコードベースが少ない、つまり***読みやすい***
3. 結構ホットな話題の割に意外と多言語に移植している人はいない（自分が観測した限り [rust に移植した日本人１名](https://github.com/HikaruEgashira/vulnhuntrs) と 自分の[goの移植](https://github.com/YmBIgo/vulnhuntr-go) のみ。Githubって、地域で検索に制限かけているんですかね？）

ここまで聞いて心踊らなかったらすみません。
ただ、もし少しでも心躍って「我こそも次の Cursor を作りたい！」と思った方は、是非とも続きを読んでくださると嬉しいです。

※：自分もセキュリティは大学の時にちょっと習って、業務で脆弱性を見つけて嬉々としているだけのフロントエンジニアなので、なんか間違いがあったらマサカリくださいね。

## LLM セキュリティ、その不可能の歴史と転換点 

まずは、簡単に LLM のセキュリティがなぜこれまでそこまで上手くいかなかったかを説明します。

まず、ChatGPT の o3-mini がセキュリティにおいて、小学生〜大学院生までだとどのレベルにいるかご存知ですか？

実は、こちらは [OpenAI](https://cdn.openai.com/o3-mini-system-card-feb10.pdf) も資料を出していて、どうやら高校生レベルのセキュリティ問題なら61%できるが、大学生・大学院生レベルだと 21% にとどまるという結果だったそうです。

この結果からも、o3-mini の紹介ページの一部では、セキュリティはレベルが low であると書いてあります。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/92647/70ecae02-305d-49ef-b642-a811797f568c.png)

https://openai.com/index/o3-mini-system-card/

しかし、LLM の限界はこんなものではありませんでした。
時間は前後しますが、2024/10/19 に ProtectAI というセキュリティ企業は、「世界で初めての AI駆動の0day探索ソフトウェア VulnHuntr」を公開しました。

https://protectai.com/threat-research/vulnhuntr-first-0-day-vulnerabilities

これの何が画期的だったかといえば、
今までの方法、つまり静的解析では false positives が多すぎて、かといって LLM だけを使いすべてのファイルを詰め込むのはコンテキストサイズの限界や文脈を把握しきれない問題がありました。
そこで [VulnHuntr](https://github.com/protectai/vulnhuntr/tree/main?tab=readme-ov-file) は まず http が最初に通るエンドポイントを特定し、そこから静的解析を通じて取得した文脈を LLM に渡し LLM に候補を絞らせることで問題を解決しました。

結果、[VulnHuntr](https://github.com/protectai/vulnhuntr/tree/main?tab=readme-ov-file) は Github 上の スターが1000以上あるプロジェクトで11もの脆弱性を発見しました。

[VulnHuntr](https://github.com/protectai/vulnhuntr/tree/main?tab=readme-ov-file) は、１つのLLMを使ったセキュリティを試みとしてはとても有意義なものになったのです。

## 実際に作ってみた

ただ、[VulnHuntr](https://github.com/protectai/vulnhuntr/tree/main?tab=readme-ov-file) は POC で python のものしか作られていません。
なので、普段業務で少し触れている golang で実装することにしました（後になって rust実装で java, js, ts, go, python, rust に対応している[神](https://github.com/HikaruEgashira/vulnhuntrs)を見つけました）。

成果物は、

https://github.com/YmBIgo/vulnhuntr-go

です。
折角なので、意気込みのある皆様にも別言語での移植を作ってもらいたいので、自分の実装の説明をしたいと思います。
ファイル構成は、

https://github.com/YmBIgo/vulnhuntr-go/tree/main/src

以下のみになっていて、VSCode拡張を作れることを夢見て、TypeScript で書きました。なぜ TypeScript で GoLang のコード解析ができるかというと、gopls という golang の language server があるからです。

https://github.com/golang/tools/tree/master/gopls

この gopls は brew で入れておけば、以下のようなコマンドで関数の定義箇所を特定できます。

```bash
gopls definition /some-path/main.go#100
```

100の部分はオフセットで、ファイルの中の指定したい関数名のIndexを示します（ここら辺は gopls を手元にダウンロードして、`gopls help definition` などを叩いてみてみてください）。

これらはあくまでも CLI なので、TypeScript なら child_process を使って実行と結果取得が可能になっています。

---

さて、話が少し脱線しましたが、src 配下のファイルはそれぞれ以下のような役割を持っています。

- api.ts　：　Claude のAPIを呼び出すためのもの
- goplsSymbolExtractor.ts　：　gopls でファイルの定義を取得するためのもの
- prompts.ts　：　プロンプトが書いてあるもの
- rootFinder.ts　：　最初に探索するエンドポイントを取得するもの
- index.ts　：　実際に探索をするもの

これらは python 版の Vulnhunhr を移植したものなので、

https://github.com/protectai/vulnhuntr/tree/main/vulnhuntr

の実装も参考になると思います。

自分の場合は、おおむね以下の手順でコードを書きました。

1. goplsSymbolExtractor.ts で、gopls で情報が正しく取れるかの確認
2. rootFinder.ts で、最初のエンドポイントを取得しつつ、そこから goplsSymbolExtractor.ts を使い正しくファイル定義を取得できるか確認
3. api.ts の実装
4. prompts.ts の実装（VulnHuntr のものを ChatGPT に聞いて golang 版にしてもらいました）
5. index.ts でメインロジックを書く

それぞれのはまりポイントをかくと、

１ goplsSymbolExtractor.ts
   A. gopls に渡すファイルの関数の位置がずれて困る問題
   B. gopls の definition か references か implementation どれを使えばいいか分からない問題

```typescript: goplsSymbolExtractor.ts
        const findReferencesCommand = `cd ${this.rootPath};
${this.goplsPath} definition ${this.filePath}:${this.fileRow}:${this.filePos}`
        // console.log(findReferencesCommand)
        let stdout, stderr;
        try {
            const std = await promisifyExec(findReferencesCommand)
            stdout = std.stdout
            stderr = std.stderr
        } catch(e) {
            console.warn(e)
            return ["", ""]
        }
```

２ rootFinder.ts
   A. chi と http に対応しているのですが、各ルーターに対応した正規表現を書くのが少しだるい
   B. goplsSymbolExtractor.ts を使って探索するときの、関数の位置を取得する問題

```typescript: rootFinder.ts
    initRegex() {
        const patterns = [
            // chi 1行の関数の場合 ()
            /Use\(([a-zA-Z0-9.]+(\([a-zA-Z0-9.]+\))*)\)/g,
            /Mount\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\(([a-zA-Z0-9.]+)\)\s*\)/g,
            /Handle\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\(([a-zA-Z0-9.]+)\)\s*\)/g,
            /HandleFunc\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\(([a-zA-Z0-9.]+)\)\s*\)/g,
            /Connect\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\(([a-zA-Z0-9.]+)\)\s*\)/g,
            /Delete\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\(([a-zA-Z0-9.]+)\)\s*\)/g,
            /Get\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\(([a-zA-Z0-9.]+)\)\s*\)/g,
            /Head\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\(([a-zA-Z0-9.]+)\)\s*\)/g,
            /Options\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\(([a-zA-Z0-9.]+)\)\s*\)/g,
            /Patch\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\(([a-zA-Z0-9.]+)\)\s*\)/g,
            /Post\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\(([a-zA-Z0-9.]+)\)\s*\)/g,
            /Put\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\(([a-zA-Z0-9.]+)\)\s*\)/g,
            /Trace\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\(([a-zA-Z0-9.]+)\)\s*\)/g,
            /NotFound\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\(([a-zA-Z0-9.]+)\)\s*\)/g,
            // http 1行の関数の場合
            /Mount\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\s*\)/g,
            /Handle\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\s*\)/g,
            /HandleFunc\(\s*"[a-zA-Z0-9\/_-\{\}#]+"\s*,\s*([a-zA-Z0-9.]+)\s*\)/g,
            /Connect\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\s*\)/g,
            /Delete\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\s*\)/g,
            /Get\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\s*\)/g,
            /Head\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\s*\)/g,
            /Options\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\s*\)/g,
            /Patch\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\s*\)/g,
            /Post\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\s*\)/g,
            /Put\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\s*\)/g,
            /Trace\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\s*\)/g,
            /NotFound\(\s*"[a-zA-Z0-9\/_-\{\}]+"\s*,\s*([a-zA-Z0-9.]+)\s*\)/g,
        ]
```

３ api.ts
   A. scratchpad という claude の機能くらい？本当は READMEもシステムコマンドとして読み込ませたい
４ index.ts
   A. 1分あたりのAPIの制限をかい潜るための sleep
   B. Claude に API投げる時に XML にする変換（まぁそんなハマらないかも）
   C. 深さ10で検索する全体的なロジックの流れ

```typescript: index.ts
                    const vulnerabilitiesXml = generateXml(
                        fileContent,
                        fixedFilePath,
                        prevAnalysis,
                        codeDefinitions,
                        VULN_SPECIFIC_BYPASSES_AND_PROMPTS[vul].bypasses.join("\n")
                    )
                    try {
                        await new Promise((resolve) => setTimeout(resolve, 5000))
                        console.log("sleeping 5 seconds")
                        const res = await llm.chat(vulnerabilitiesXml)
                        secondReportJson = parseJsonRes(JSON.parse(res))
                        console.log(">>>> second report iteration ... ", i, currentFile)
                    } catch(e) {
                        console.error(e)
                        secondReportJson = {} as AnalysisResult
                    }
                    if (!secondReportJson.context_code || !secondReportJson.context_code.length) {
                        break
                    }
```

みたいな感じになりますが、コードもあるんでよかったら内容を見てみてください。

https://github.com/YmBIgo/vulnhuntr-go/tree/main/src

この手順は、他の言語に転用するときも使えます。
例えば ruby なら [ruby-lsp](https://github.com/Shopify/ruby-lsp) という gem の language server がありますが、ruby ファイルを作ってその中で ruby-lsp を実行させることで、コマンドラインライクに使うことができると思います。
まずは、１つずつ着実に実装していけば、似たようなものはできると思います。

## 今後の展望

JavaScript・TypeScript を主戦場にしているので、そこら辺で移植版を書ければいいなと思います。
また、ruby も日本ではまだ使われていると思うので、ruby 版を作ってみたいですね。

ここまで読んでいただきありがとうございます。この記事を通じて、セキュリティLLMに興味を持つ人が１人でも増えれば幸いです。
ありがとうございます。

## 参考文献

https://zenn.dev/sqer/articles/0ab529884a197f

https://github.com/protectai/vulnhuntr/tree/main

https://protectai.com/threat-research/vulnhuntr-first-0-day-vulnerabilities