# Pickleスキャン

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

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

## Pickle とは？

[公式ドキュメント](https://docs.python.org/3/library/pickle.html)より:

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

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

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

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

以下は小さな例です：

```python
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ファイルが作成され、ターミナルに以下の手順が表示されます。

```python
    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](https://docs.python.org/3/library/pickletools.html)モジュールはpickleの分析に非常に有用である、ということだけ知っておいてください。このモジュールを使用すると、***実行せずに***ファイル内の指示を読み取ることができます。

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

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

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

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

```python
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`が生成されます。ファイルの内容を確認すると:

```bash

# 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に何が問題なのか、と思われるかもしれません。

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

```python
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](https://github.com/trailofbits/fickling) ライブラリを使用しています。これにより、`exec` 関数を通じて文字列に含まれるコードを実行するpickle指示を追加できます。pickle内で関数やクラスを定義できないという制約を回避する方法は、文字列として保存されたPythonコードに対してexecを実行することです。

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

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

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

```bash
# 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×.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘
```

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

```python
# ...
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キーでコミットに署名する](./security-gpg)ことができます。これによりファイルが安全であることは**保証されません**が、ファイルの由来は保証されます。

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

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

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

E.g.:

```python
from transformers import AutoModel

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

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

- [MsgPack](https://msgpack.org/index.html)
- [Protobuf](https://developers.google.com/protocol-buffers)
- [Cap'n'proto](https://capnproto.org/)
- [Avro](https://avro.apache.org/)
- [safetensors](https://github.com/huggingface/safetensors)

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

### `torch.load/save`の改善

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

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

#### 現在の状況

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

- ClamAVスキャン
- Pickle Import スキャン

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

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

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

このデータは、[`pickletools.genops`](https://docs.python.org/3/library/pickletools.html#pickletools.genops) によって取得できます。これにより、潜在的に危険なコードを実行せずにファイルを読み取ることができます。

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

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

#### 考えられる解決策

カスタムの[Unpickler](https://docs.python.org/3/library/pickle.html#pickle.Unpickler)を作成する、という考え方もあります。[この例](https://github.com/facebookresearch/CrypTen/blob/main/crypten/common/serial.py)のようなものです。しかし、[この洗練されたエクスプロイト](https://ctftime.org/writeup/16723)で確認できるように、この方法は機能しません。

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

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

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

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

## さらに詳しく

[pickle - Pythonオブジェクトのシリアライズ - Python 3.10.6 ドキュメント](https://docs.python.org/3/library/pickle.html#what-can-be-pickled-and-unpickled)

[危険な Pickles - 悪意のある Python シリアライズ](https://intoli.com/blog/dangerous-pickles/)

[GitHub - trailofbits/fickling: Pythonのピクルスデコンパイラおよび静的解析ツール](https://github.com/trailofbits/fickling)

[Python pickleの悪用](https://davidhamann.de/2020/04/05/exploiting-python-pickle/)

[cpython/pickletools.py at 3.10 · python/cpython](https://github.com/python/cpython/blob/3.10/Lib/pickletools.py)

[cpython/pickle.py at 3.10 · python/cpython](https://github.com/python/cpython/blob/3.10/Lib/pickle.py)

[CrypTen/serial.py at main · facebookresearch/CrypTen](https://github.com/facebookresearch/CrypTen/blob/main/crypten/common/serial.py)

[CTFtime.org / Balsn CTF 2019 / pyshv1 / 解説](https://ctftime.org/writeup/16723)

[Pythonのpickleモジュールの再整備](https://github.com/moreati/pickle-fuzz)