GPTに基づくネットワークセキュリティツールの設計

最近も GPT puls が開通し、GPT4 の効果が確かに 3.5 よりも強いと感じています。単に文脈能力や文字数制限だけでなく、論理能力もかなり向上しています。ついに生産性の解放に使えるようになったので、この記事が生まれました。専門的な開発ではなく、コードも稚拙なので、大御所の方々は軽くスルーしてください。

01#

ご存知の通り、安全職の仕事の大部分は報告書の作成です。しばしば、脆弱性を発見するのに 5 分かかるのに対し、報告書を書くのには 15 分かかります。報告書は特定の形式と特定の文法で書かなければならず、本当に苦痛な作業です。そこで、AI に私たちを苦痛から解放してもらう必要があります。

設計思路#

UI インターフェースを持つ
ファイル処理が容易
AI に基づいて言語記述の問題を処理
人工的に結果に干渉できる

まず、UI インターフェースを作成します。迅速な開発のために Python を選び、軽量開発のために Web を選びました。

flask を使って始めます。

app = Flask(__name__, template_folder='templates', static_folder='static')

@app.route('/', methods=['GET', 'POST'])
def index():
    if request.method == 'POST':
        vuln_name = request.form.get('vuln_name', default="SQL注入")
        vuln_point = request.form.get('vuln_point', default="www.google.com")
        beizhu = request.form.get('beizhu', default="")
        language = request.form.get('language')
        start = time.time()
        report = generate_report(vuln_name, vuln_point, beizhu, language)
        end = time.time()
        times = end - start

        return render_template('index.html', report=report, times=times)
    else:
        return render_template('index.html')

デフォルトのルートを作成し、テンプレート index.html を読み込みました。
index.html の主体はフォームです。

フォーム：

脆弱性名称（例：SQL 注入）
脆弱性ポイント（例：www.test.com/id=111）
备注（例えば、報告を詳細にする、xx 個の POC を提供するなどの追加要求）

次に、このページを簡単に実装するためにもう一つの HTML を書きます。

<!DOCTYPE html>
<html>

<head>
  <meta charset="UTF-8">
  <title>脆弱性報告生成器</title>
  <style>
    body {
      font-family: Arial, sans-serif;
      margin: 0;
      padding: 0;
    }

    header {
      background-color: #005293;
      color: #fff;
      padding: 10px;
    }

    h1 {
      margin: 0;
    }

    form {
      margin: 20px;
    }

    label {
      display: block;
      margin-bottom: 10px;
    }

    input[type="text"] {
      width: 100%;
      padding: 10px;
      border: 1px solid #ccc;
      border-radius: 5px;
      margin-bottom: 20px;
      font-size: 16px;
    }

    input[type="submit"] {
      background-color: #005293;
      color: #fff;
      border: none;
      padding: 10px;
      border-radius: 5px;
      cursor: pointer;
      font-size: 16px;
    }

    #report {
      margin: 20px;
      border: 1px solid #ccc;
      padding: 10px;
      border-radius: 5px;
      font-size: 16px;
    }
  </style>

  <link rel="stylesheet"
    href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.7.0/styles/panda-syntax-dark.min.css">

  <script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.7.0/highlight.min.js"
    integrity="sha512-bgHRAiTjGrzHzLyKOnpFvaEpGzJet3z4tZnXGjpsCcqOnAH6VGUx9frc5bcIhKTVLEiCO6vEhNAgx5jtLUYrfA=="
    crossorigin="anonymous" referrerpolicy="no-referrer"></script>

  <script>hljs.initHighlightingOnLoad();</script>
  <script src="https://cdn.jsdelivr.net/npm/[email protected]/base64.min.js"></script>
</head>

<body>
  <header>
    <h1>脆弱性報告生成器</h1>

  </header>

  <form method="POST" action="/">

    <select name="language">
      <option value="中文">中文</option>
      <option value="英文">English</option>
    </select>

    <label>脆弱性名称：</label>
    <input type="text" name="vuln_name" placeholder="脆弱性名称を入力してください">
    <label>脆弱性ポイント：</label>
    <input type="text" name="vuln_point" placeholder="脆弱性ポイントを入力してください">
    <label>報告に対する追加要求：</label>
    <input type="text" name="beizhu" placeholder="報告の追加要求">

    <input type="submit" value="脆弱性報告を生成"> <button id="copy-btn" onclick="copyReport()" type="button">コピー</button><!--<button id="copy-btn" onclick="saveReport()" type="button">ファイルを生成</button>-->
    実行時間：{{ times }} 秒
  </form>


  {% if report %}

  <pre id="report" style="white-space: pre-wrap;overflow-wrap: break-word;">
        <code>
      {{ report }}
      </code>
    </pre>


  <br>

  {% if report %}

  {% endif %}

  <script>
    function copyReport() {
      var reportText = document.querySelector("#report code").innerText;
      navigator.clipboard.writeText(reportText)
        .then(function () {
          alert("脆弱性報告がクリップボードにコピーされました！");
        })
        .catch(function (error) {
          alert("コピーに失敗しました。手動でコピーしてください。");
        });
    }
  </script>

  {% endif %}
</body>

</html>

ページは大体こんな感じで、いくつかの小さな詳細を改善します。

次にファイル処理です。人が快適に読み、ブラウザのテキストエディタで解析できるファイルやエンコーディング形式はありますか？そうです、Markdown が最初に思い浮かびます。そこで、index.html にコードハイライトを追加し、生成された値をその中に直接埋め込みました。

最終的な実現効果：

次に人工的な調整ですが、実際にはすでに備考オプションで完了しているので、これ以上は述べません。

最後に最も重要なこと

chatGPT API と接続し、情報を取得します。
ここで、GitHub のプロジェクトをお勧めします。chatGPT の開発資料はここで確認できます：https://github.com/easychen/openai-gpt-dev-notes-for-cn-developer

次に、ロジックコードを書きます。

def generate_report(vuln_name, vuln_point, beihzu="", language=""):
    # GPT-3の入力を構築
    prompt = f"" #ここに調整したプロンプトを記入

    #api_base = {"socks5":proxy}
    api_base = {"http": proxy, "https": proxy}

    # OpenAI APIを呼び出して脆弱性報告を生成
    headers = {
        # json=を渡すときにすでに追加されますが、data=を渡すときには追加されません
        'Content-Type': 'application/json',
        'Authorization': "Bearer "+api_key,
    }

    json_data = {
        'model': 'gpt-3.5-turbo',
        'messages': [
            {
                'role': 'user',
                'content': prompt,
            },
        ],
    }
    # プロキシを設定
    # 時間を計算

    response = requests.post('https://api.openai.com/v1/chat/completions',
                             headers=headers, json=json_data, proxies=api_base)

    # OpenAI APIの応答を処理
    if response.status_code == 200:
        text = response.text
        text = json.loads(text)
        text = text['choices'][0]['message']['content']
        report = f"{text}"

    else:
        report = "脆弱性報告の生成に失敗しました。入力を確認して再試行してください。ステータスコード："+str(response.status_code)

    return report

ここでは公式の方法を使用せず、プロキシのために HTTP リクエストインターフェースを使用しました。

大体のロジックはこんな感じで、ついでにファイル保存機能も書きます。

def outfile():
    # POSTで送信されたbase64エンコードされたデータを取得
    date=request.form.get('reportText')
    
    # データをbase64デコード
    date=base64.b64decode(date,).decode('utf-8')
    # データをmdファイルに書き込み、ファイル名は現在の時間
    fname=time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
    with open(fname+'.md','w') as f:
        f.write(date)
    return "ok"

全体的にはこんな感じで、最終的な効果図

最後にまとめると、プロジェクトは
https://github.com/shiyeshu/GPTreport

大御所の方々は気に入ったらスターをください。

02#

実はこれが本題ですが、会社に奇妙な要求がありました。四月末までに特許を登録して会社に献上する必要があります。考えを絞って、この献上を準備するしかありません。

大まかなアイデアを述べます。

AI のコード監査ツールを作成することは、前のものよりも複雑で実用的です。

まず、AI は全ソースコードを全文解析することは不可能であり、小さな部分の解析しかできません。
そこで、手段を使って事前処理を行います。いくつかのソフトウェアの脆弱性マッチングルールを使用して事前スキャンを行うことで、特定の行のコードにリスクがあることを確認できます。その後、その行のコードに対して 2 つの操作を行います。関数内にあるかどうかを判断します。関数であれば、関数全体を分析し、そうでなければ、ある閾値内で前後の文脈を連動して分析します。
さらに高度な方法として、変数の代入操作を追跡することもできます。これで、基本的にプログラム全体の監査ロジックが完成します。
実際には、遊べる方向性はまだたくさんありますが、ここではこれ以上は述べません。
コードはすでに書き終え、今手元に持っています。もし四月が終わっても会社が提出を催促しなければ、オープンソースにします。
少しだけ画像を公開します。

最後に
私はあまり記事を書かないので文才がなく、言葉が混乱しています。皆さん、どうかご容赦ください。この記事で言及されているいくつかのことは比較的ホットな話題ですので、皆さんは純粋に技術的な視点で見て、オープンソースを受け入れつつ、ただのタダ乗りの泥棒にならないようにしてください。