Hub documentation

Pickleスキャン

Hub

API docs

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

Pickleスキャン

Pickleは機械学習で広く使用されているシリアライズ形式です。特に、PyTorchのモデル重みのデフォルト形式として知られています。

pickle ファイルを読み込む際に、危険な任意のコード実行攻撃が実行される可能性があります。信頼できるユーザーや組織からモデルを読み込み、署名付きコミットに依存する、および/またはfrom_tf=Trueの自動変換メカニズムを使用して TF または Jax 形式からモデルを読み込むことをお勧めします。また、この問題を軽減するために、pickle 化されたファイル内のインポートのリストを Hub 上で直接表示/「審査」しています。最後に、safetensorsと呼ばれる新しいシンプルな重みのシリアライズ形式の実験を行っています。

Pickle とは？

公式ドキュメントより:

pickleモジュールは、Pythonのオブジェクト構造をシリアライズおよびデシリアライズするためのバイナリプロトコルを実装しています。

これは、pickleがシリアライズプロトコルであり、パーティ間でデータを効率的に共有するために使用されるものという意味です。

ピクルスとは、ピクルス化の過程で生成されたバイナリファイルを指します。

pickle の中核は、基本的に命令（オペコード）のスタックです。おそらくお気づきのことと思いますが、人間が読める形式ではありません。オペコードは pickle 化（シリアライズ）時に生成され、unpickle 化（デシリアライズ）時に順次読み取られます。オペコードに基づいて、特定のアクションが実行されます。

以下は小さな例です：

import pickle
import pickletools

var = "data I want to share with a friend"

# store the pickle data in a file named 'payload.pkl'
with open('payload.pkl', 'wb') as f:
    pickle.dump(var, f)

# disassemble the pickle
# and print the instructions to the command line
with open('payload.pkl', 'rb') as f:
    pickletools.dis(f)

これを実行すると、pickleファイルが作成され、ターミナルに以下の手順が表示されます。

    0: \x80 PROTO      4
    2: \x95 FRAME      48
   11: \x8c SHORT_BINUNICODE 'data I want to share with a friend'
   57: \x94 MEMOIZE    (as 0)
   58: .    STOP
highest protocol among opcodes = 4

今は指示の詳細についてあまり気にしなくて大丈夫です。pickletoolsモジュールはpickleの分析に非常に有用です。このモジュールを使用すると、実行せずにファイル内の指示を読み取ることができます。

Pickleは単なるシリアライズプロトコルではなく、デシリアライズ時にPythonコードを実行できる柔軟性をユーザーに提供します。あまり良い印象ではないですよね？

なぜ危険なのでしょうか？

上記で述べたように、pickleのデシリアライズはコードの実行を意味します。ただし、これにはいくつかの制限があります。トップレベルのモジュールから関数やクラスを参照することは可能ですが、pickleファイル自体にそれらを埋め込むことはできません。

もう一度ゼロから考え直す:

import pickle
import pickletools

class Data:
    def __init__(self, important_stuff: str):
        self.important_stuff = important_stuff

d = Data("42")

with open('payload.pkl', 'wb') as f:
    pickle.dump(d, f)

このスクリプトを実行すると、再びpayload.pklが生成されます。ファイルの内容を確認すると:


# cat payload.pkl
__main__Data)}important_stuff42sb.%

# hexyl payload.pkl
┌────────┬─────────────────────────┬─────────────────────────┬────────┬────────┐
│00000000│ 80 04 95 33 00 00 00 00 ┊ 00 00 00 8c 08 5f 5f 6d │×•×30000┊000×•__m│
│00000010│ 61 69 6e 5f 5f 94 8c 04 ┊ 44 61 74 61 94 93 94 29 │ain__××•┊Data×××)│
│00000020│ 81 94 7d 94 8c 0f 69 6d ┊ 70 6f 72 74 61 6e 74 5f │××}××•im┊portant_│
│00000030│ 73 74 75 66 66 94 8c 02 ┊ 34 32 94 73 62 2e       │stuff××•┊42×sb.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘

中身を見ると、オPCODEと関連するデータがいくつかあるだけであることがわかります。pickleに何が問題なのか、と思われるかもしれません。

別のことを試してみましょう:

from fickling.pickle import Pickled
import pickle

# Create a malicious pickle
data = "my friend needs to know this"

pickle_bin = pickle.dumps(data)

p = Pickled.load(pickle_bin)

p.insert_python_exec('print("you\'ve been pwned !")')

with open('payload.pkl', 'wb') as f:
    p.dump(f)

# innocently unpickle and get your friend's data
with open('payload.pkl', 'rb') as f:
    data = pickle.load(f)
    print(data)

ここでは簡便さのため、fickling ライブラリを使用しています。これにより、exec 関数を通じて文字列に含まれるコードを実行するpickle指示を追加できます。pickle内で関数やクラスを定義できないという制約を回避する方法は、文字列として保存されたPythonコードに対してexecを実行することです。

これを実行すると、payload.pkl が作成され、以下が出力されます。

you've been pwned !
my friend needs to know this

pickleファイルの内容を確認すると、次のようになります。

# cat payload.pkl
c__builtin__
exec
(Vprint("you've been pwned !")
tR my friend needs to know this.%

# hexyl payload.pkl
┌────────┬─────────────────────────┬─────────────────────────┬────────┬────────┐
│00000000│ 63 5f 5f 62 75 69 6c 74 ┊ 69 6e 5f 5f 0a 65 78 65 │c__built┊in___exe│
│00000010│ 63 0a 28 56 70 72 69 6e ┊ 74 28 22 79 6f 75 27 76 │c_(Vprin┊t("you'v│
│00000020│ 65 20 62 65 65 6e 20 70 ┊ 77 6e 65 64 20 21 22 29 │e been p┊wned !")│
│00000030│ 0a 74 52 80 04 95 20 00 ┊ 00 00 00 00 00 00 8c 1c │_tR×•× 0┊000000×•│
│00000040│ 6d 79 20 66 72 69 65 6e ┊ 64 20 6e 65 65 64 73 20 │my frien┊d needs │
│00000050│ 74 6f 20 6b 6e 6f 77 20 ┊ 74 68 69 73 94 2e       │to know ┊this×.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘

基本的には、アンピクルを行う際に以下のような処理が行われます。

# ...
opcodes_stack = [exec_func, "malicious argument", "REDUCE"]
opcode = stack.pop()
if opcode == "REDUCE":
    arg = opcodes_stack.pop()
    callable = opcodes_stack.pop()
    opcodes_stack.append(callable(arg))
# ...

脅威となる命令は、STACK_GLOBAL、GLOBAL、REDUCEです。

REDUCE は、アンピクルラーに指定された引数で関数を実行することを指示し、*GLOBAL 命令はアンピクルラーにimport を実行することを指示します。

まとめると、pickleが危険な理由は以下の通りです：

Pythonモジュールをインポートする際、任意のコードが実行される可能性があります
evalやexecなどの組み込み関数をインポートでき、これらは任意のコードを実行するために使用できます
オブジェクトをインスタンス化する際、コンストラクタが呼び出される場合があります

そのため、pickle を使用しているほとんどのドキュメントでは、信頼できないソースからのデータをアンピクルしないよう注意喚起されています。

緩和策

pickleは使用しないでください

Luc、的確な助言ですね。ただし、pickle は広く使用されており、近々姿を消すことはありません。すべての関係者が納得できる新しいフォーマットを見つけ、変更を推進するには、ある程度の時間がかかるでしょう。

では、現時点で何ができるでしょうか？

信頼できるユーザーや組織からファイルをロードします

Hubでは、GPGキーでコミットに署名することができます。これによりファイルが安全であることは保証されませんが、ファイルの由来は保証されます。

ユーザーAを信頼しており、Hub上のファイルを含むコミットがユーザーAのGPGキーによって署名されている場合、そのファイルを信頼して問題ないと判断してほぼ安全です。

TFまたはFlaxからモデルの重みをロードする

TensorFlowおよびFlaxのチェックポイントは影響を受けず、from_pretrainedメソッドのfrom_tfおよびfrom_flaxキーワード引数を使用してPyTorchアーキテクチャ内で読み込むことで、この問題を回避できます。

E.g.:

from transformers import AutoModel

model = AutoModel.from_pretrained("google-bert/bert-base-cased", from_flax=True)

独自のシリアライズ形式を使用する

この最後の形式であるsafetensorsは、現在開発・実験を進めているシンプルなシリアライズ形式です！可能であれば、ご協力やコントリビューションをお願いします🔥。

torch.load/saveの改善

PyTorchでは、デフォルトで*.ptファイルから重みのみを安全に読み込む方法に関するオープンな議論が進行中です。ぜひそちらに参加してください！

Hubのセキュリティスキャナー

現在の状況

セキュリティスキャナーを作成しました。このスキャナーは、Hubにプッシュされたすべてのファイルをスキャンし、セキュリティチェックを実行します。執筆時点では、2種類のスキャンを実行しています。

ClamAVスキャン
Pickle Import スキャン

ClamAVのスキャンでは、ファイルはオープンソースのアンチウイルスソフトであるClamAVでスキャンされます。これにより多くの危険なファイルが検出されますが、pickleの悪用には対応していません。

Pickle Import スキャンを実装しました。これは、pickle ファイル内で参照されているインポートのリストを抽出します。pytorch_model.bin またはその他の pickle 化されたファイルをアップロードするたびに、このスキャンが実行されます。

Hub 上では、インポートを含む各ファイルの横にインポートのリストが表示されます。怪しいインポートがある場合は、ハイライト表示されます。

このデータは、pickletools.genops によって取得できます。これにより、潜在的に危険なコードを実行せずにファイルを読み取ることができます。

ファイルのアンピクル化時に、*GLOBALによってインポートされた潜在的に危険な関数に対してREDUCEが実行されるかどうかを判断できるのは、この仕組みのおかげです。

免責事項: これは100%確実ではありません。ユーザーとして、何かが安全かどうかを確認する責任はあなたにあります。私たちはPythonパッケージの安全性を能動的に監査しておらず、安全/不安全なインポートのリストは最善の努力で維持されています。何か安全でないと思われるものがある場合は、website at huggingface.co 宛にメールをお送りください。そうすれば、私たちはそれを安全でないとしてフラグ付けします。

考えられる解決策

カスタムのUnpicklerを作成する、という考え方もあります。この例のようなものです。しかし、この洗練されたエクスプロイトで確認できるように、この方法は機能しません。

幸いなことに、evalのインポートの痕跡は常に残るため、オPCODEを直接読み取れば、悪意のある使用を検出できるはずです。

現在提案している解決策は、.gitignoreに似たファイルを作成し、インポート用にするというものです。

このファイルは、ホワイトリストにないインポートが存在する場合にpytorch_model.binファイルが危険と判定される、インポートのホワイトリストです。

例えば、numpy.*のようなシンプルな1行で、すべてのnumpyサブモジュールを許可できるような、正規表現風のフォーマットを想像できます。

さらに詳しく

pickle - Pythonオブジェクトのシリアライズ - Python 3.10.6 ドキュメント

危険な Pickles - 悪意のある Python シリアライズ

GitHub - trailofbits/fickling: Pythonのピクルスデコンパイラおよび静的解析ツール

Python pickleの悪用

cpython/pickletools.py at 3.10 · python/cpython

cpython/pickle.py at 3.10 · python/cpython

CrypTen/serial.py at main · facebookresearch/CrypTen

CTFtime.org / Balsn CTF 2019 / pyshv1 / 解説

Pythonのpickleモジュールの再整備

Update on GitHub

←マルウェアスキャンシークレットスキャン→