# Pickle 스캔

Pickle은 ML에서 널리 사용되는 직렬화 형식입니다. 특히 PyTorch 모델 가중치의 기본 형식으로 사용됩니다.

pickle 파일을 로드할 때 위험한 임의 코드 실행 공격이 발생할 수 있습니다. 신뢰할 수 있는 사용자 및 조직에서 모델을 로드하고, 서명된 커밋에 의존하며, `from_tf=True` 자동 변환 메커니즘을 사용하여 TF 또는 Jax 형식의 모델을 로드하는 것을 권장합니다. 또한 Hub에서 pickle 파일의 import 목록을 직접 표시/검토하여 이 문제를 완화합니다. 마지막으로, [`safetensors`](https://github.com/huggingface/safetensors)라고 불리는 가중치를 위한 새로운 단순 직렬화 형식을 실험하고 있습니다.

## pickle란 무엇인가요?

[공식 문서](https://docs.python.org/3/library/pickle.html)에서:

> `pickle` 모듈은 Python 객체 구조를 직렬화 및 역직렬화하기 위한 바이너리 프로토콜을 구현합니다.

이는 pickle이 직렬화 프로토콜이라는 것을 의미하며, 여러 당사자 간에 데이터를 효율적으로 공유하는 데 사용하는 것입니다.

pickle는 pickling 과정에서 생성된 바이너리 파일을 가리킵니다.

핵심적으로, pickle은 기본적으로 명령어나 opcode의 스택입니다. 짐작하셨을 것처럼, 이는 사람이 읽기 어렵습니다. opcode는 pickling 시 생성되며, unpickling 시 순차적으로 읽힙니다. opcode에 따라 특정 동작이 실행됩니다.

다음은 간단한 예시입니다:

```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의 역직렬화(de-serialization)는 코드가 실행될 수 있음을 의미합니다. 하지만 여기에는 몇 가지 제한 사항이 있습니다. 최상위 모듈의 함수와 클래스만 참조할 수 있으며, 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.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘
```

여기에는 opcode 몇 개와 관련 데이터가 있을 뿐, 내용이 많지 않음을 확인할 수 있습니다. 그렇다면 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×.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘
```

기본적으로 unpickle을 수행할 때 발생하는 과정은 다음과 같습니다:

```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`는 unpickler에게 제공된 인수를 사용하여 함수를 실행하도록 지시하는 것이며, `*GLOBAL` 명령은 unpickler에게 `import`를 수행하도록 지시합니다.

요약하자면, pickle이 위험한 이유는 다음과 같습니다:

- python 모듈을 가져올 때 임의의 코드가 실행될 수 있습니다
- `eval` 또는 `exec`과 같은 내장 함수를 가져와 임의의 코드를 실행할 수 있습니다
- __init__ 메서드는 객체를 인스턴스화할 때 생성자로 호출될 수 있습니다.

pickle을 사용하는 대부분의 문서에서 신뢰할 수 없는 출처의 데이터를 언픽클링하지 말라고 명시하는 이유가 바로 이것입니다.

## 완화 전략

***pickle을 사용하지 마세요***

좋은 조언입니다, Luc. 하지만 pickle는 여전히 광범위하게 사용되고 있으며 당분간 사라지지 않을 것입니다. 모두가 만족할 수 있는 새로운 포맷을 찾아 변경을 시작하는 데는 상당한 시간이 소요될 것입니다.

지금 당장 무엇을 할 수 있을까요?

### 신뢰할 수 있는 사용자 및 조직에서 파일을 로드합니다

Hub에서는 [GPG 키로 커밋에 서명](./security-gpg)할 수 있습니다. 이는 파일이 안전하다는 것을 **보장하지는** 않지만, 파일의 출처는 보장합니다.

사용자 A와 Hub에 있는 해당 파일이 포함된 커밋이 사용자 A의 GPG 키로 서명되었다는 것을 알고 신뢰한다면, 그 파일을 신뢰할 수 있다고 가정하는 것이 상당히 안전합니다.

### TF 또는 Flax에서 모델 가중치 로드

TensorFlow 및 Flax 체크포인트에는 영향이 없으며, `from_pretrained` 메서드의 `from_tf` 및 `from_flax` kwargs를 사용하여 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에 push된 모든 파일을 스캔하여 보안 검사를 실행하는 보안 스캐너를 만들었습니다. 작성 시점 기준으로 두 가지 유형의 스캔을 실행합니다:

- ClamAV 스캔
- Pickle Import 스캔

ClamAV 스캔의 경우, 파일은 오픈소스 안티바이러스 [ClamAV](https://www.clamav.net)를 통해 검사됩니다. 이는 상당수의 위험한 파일을 포괄하지만, pickle 익스플로잇은 포함하지 않습니다.

Pickle Import 스캔을 구현했으며, 이는 pickle 파일에서 참조된 import 목록을 추출합니다. `pytorch_model.bin` 또는 기타 pickled 파일을 업로드할 때마다 이 스캔이 실행됩니다.

Hub에서는 import가 포함된 각 파일 옆에 import 목록이 표시됩니다. 의심스러운 import가 있으면 강조 표시됩니다.

이 데이터는 [`pickletools.genops`](https://docs.python.org/3/library/pickletools.html#pickletools.genops) 덕분에 얻을 수 있으며, 이 함수를 사용하면 잠재적으로 위험한 코드를 실행하지 않고 파일을 읽을 수 있습니다.

`*GLOBAL`을 통해 가져온 잠재적으로 위험한 함수에 대해 `REDUCE`가 수행되는지 여부를 파악할 수 있게 해주는 것이 바로 이 부분입니다.

***면책 조항***: 이 방법은 100% 완벽하지 않습니다. 사용자가 해당 항목이 안전한지 여부를 확인하는 것이 책임입니다. 우리는 Python 패키지의 안전성을 능동적으로 감사하지 않으며, 안전/불안전한 import 목록은 최선의 노력으로 유지 관리됩니다. 안전하지 않다고 판단되는 항목이 있으면 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` import의 흔적이 항상 남아 있으므로, opcodes를 직접 읽으면 악성 사용을 감지할 수 있습니다.

제가 제안하는 현재 해결책은 `.gitignore`와 유사하지만 import를 위한 파일을 만드는 것입니다.

이 파일은 화이트리스트에 포함되지 않은 import가 있는 경우 `pytorch_model.bin` 파일을 위험한 것으로 표시하는 import의 화이트리스트입니다.

예를 들어 `numpy.*`와 같은 간단한 한 줄로 모든 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 pickling 디컴파일러 및 정적 분석기](https://github.com/trailofbits/fickling)

[Python pickles 악용하기](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 / Writeup](https://ctftime.org/writeup/16723)

[Python의 pickle 모듈 재활성화](https://github.com/moreati/pickle-fuzz)