Pickle 스캐닝

Pickle은 ML에서 널리 쓰이는 직렬화 형식이에요. 특히 PyTorch 모델 가중치의 기본 형식이기도 하죠.

출처: 문서

본문

pickle 파일을 로드할 때 위험한 임의 코드 실행 공격이 발생할 수 있어요. 신뢰하는 사용자와 조직의 모델만 로드하고, 서명된 커밋에 의존하고, TF나 Jax 형식의 모델은 from_tf=True 자동 변환 메커니즘으로 로드할 것을 권장해요. 또 Hub에서 pickle 파일에 포함된 import 목록을 직접 표시/"검토(vetting)"함으로써 이 문제를 완화하고 있어요. 마지막으로, safetensors라는 가중치용 새롭고 단순한 직렬화 형식을 실험 중이에요.

pickle이 뭔가요?

공식 문서에 따르면:

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

즉 pickle은 직렬화 프로토콜로, 당사자 간에 데이터를 효율적으로 공유하는 데 사용하는 것이에요.

pickling 과정에서 생성된 바이너리 파일을 pickle이라고 불러요.

핵심적으로 pickle은 기본적으로 명령(instruction) 또는 opcode의 스택이에요. 짐작했듯이 사람이 읽을 수는 없어요. opcode는 pickling할 때 생성되고 unpickling할 때 순차적으로 읽혀요. opcode에 따라 주어진 동작이 실행돼요.

작은 예시를 보여줄게요.

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.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘

별로 많이 들어 있지 않은 게 보이죠. 몇 개의 opcode와 관련 데이터뿐이에요. 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×.  │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘

기본적으로 이게 unpickle할 때 일어나는 일이에요.

# ...
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 모듈을 import할 때 임의 코드가 실행될 수 있음
eval이나 exec 같은 내장 함수를 import해 임의 코드를 실행하는 데 사용할 수 있음
객체를 인스턴스화할 때 생성자가 호출될 수 있음

이것이 대부분의 문서에서 "pickle을 사용할 때는 신뢰할 수 없는 소스의 데이터를 unpickle하지 말라"고 말하는 이유예요.

완화 전략

pickle을 사용하지 마세요

훌륭한 조언이지만, pickle은 널리 사용되고 머지않아 사라지지 않아요. 모두가 만족하는 새 형식을 찾고 전환을 시작하는 데는 시간이 걸릴 거예요.

그럼 지금은 무엇을 할 수 있을까요?

신뢰하는 사용자와 조직의 파일을 로드하세요

Hub에서는 GPG 키로 커밋에 서명할 수 있어요. 이것이 파일이 안전함을 보장하지는 않지만, 파일의 출처는 보장해요.

사용자 A를 알고 신뢰하고, Hub의 파일이 포함된 커밋이 사용자 A의 GPG 키로 서명되어 있다면, 그 파일을 신뢰해도 안전하다고 볼 수 있어요.

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

TensorFlow와 Flax 체크포인트는 영향을 받지 않아요. from_pretrained 메서드의 from_tf와 from_flax kwarg를 사용해 PyTorch 아키텍처 안에서 로드해 이 문제를 우회할 수 있어요.

예:

from transformers import AutoModel

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

자신의 직렬화 형식 사용

이 마지막 형식 safetensors는 우리가 현재 개발·실험 중인 단순한 직렬화 형식이에요. 가능하면 기여해 주세요 🔥.

`torch.load/save` 개선

PyTorch에 기본적으로 *.pt 파일에서 가중치만 안전하게 로드하는 방법에 대한 공개 논의가 진행 중이에요 — 참여해 주세요!

Hub의 보안 스캐너

현재 제공되는 것

Hub에 푸시된 모든 파일을 스캔하고 보안 검사를 수행하는 보안 스캐너를 만들었어요. 글을 쓰는 시점에 두 종류의 스캔을 실행해요.

ClamAV 스캔
Pickle Import 스캔

ClamAV 스캔은 오픈소스 안티바이러스 ClamAV에 파일을 통과시켜요. 위험한 파일을 상당 부분 잡아내지만 pickle 익스플로잇은 잡지 못해요.

Pickle Import 스캔을 구현했는데, pickle 파일에서 참조하는 import 목록을 추출해요. pytorch_model.bin이나 다른 pickle 파일을 업로드할 때마다 이 스캔이 실행돼요.

Hub에서 import 목록은 import를 포함하는 각 파일 옆에 표시돼요. 어떤 import가 의심스러우면 강조 표시돼요.

이 데이터는 잠재적으로 위험한 코드를 실행하지 않고 파일을 읽을 수 있는 pickletools.genops 덕분에 얻을 수 있어요.

이것으로 unpickle할 때 *GLOBAL이 import한 잠재적으로 위험한 함수에 대해 REDUCE할지 알 수 있어요.

면책 조항: 이것은 100% 완벽하지 않아요. 무언가 안전한지 확인하는 것은 사용자의 책임이에요. 우리는 Python 패키지 안전성을 적극적으로 감사하지 않고, 보유한 safe/unsafe import 목록은 best-effort로 유지돼요. 무언가 안전하지 않다고 생각되면 [email protected]로 이메일을 보내 알려주세요.

잠재적 해결책

이것 같은 커스텀 Unpickler를 만드는 것을 생각할 수 있어요. 하지만 이 정교한 익스플로잇에서 볼 수 있듯 이것은 동작하지 않아요.

고맙게도 항상 eval import의 흔적이 남으므로, opcode를 직접 읽으면 악성 사용을 잡아낼 수 있어요.

현재 제가 제안하는 해결책은 .gitignore와 비슷하지만 import용인 파일을 만드는 것이에요.

이 파일은 whitelist(허용 목록) 역할을 해서, whitelist에 없는 import가 있으면 pytorch_model.bin 파일을 위험한 것으로 표시하게 돼요.

예를 들어 numpy.* 같은 한 줄로 모든 numpy 서브모듈을 허용할 수 있는 regex 같은 형식을 상상할 수 있어요.

더 읽어보기

pickle - Python object serialization - Python 3.10.6 documentation

Dangerous Pickles - Malicious Python Serialization

GitHub - trailofbits/fickling: A Python pickling decompiler and static analyzer

Exploiting Python pickles

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 / Writeup

Rehabilitating Python's pickle module

더 알아보기 (Learn more)

pickle은 역직렬화 시 임의 코드 실행 위험 때문에 "신뢰하지 않는 소스의 데이터를 unpickle하지 말라"는 경고를 받는 형식이에요. Hub는 파일에 포함된 import 목록을 표시해 위험을 알려주고, safetensors 같은 안전한 형식을 권장해요. 신뢰하는 출처의 모델만 로드하고, 가능하면 pickle 대신 새 직렬화 형식을 쓰세요.

Pickle 스캐닝