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은 ML에서 널리 사용되는 직렬화 형식입니다. 특히 PyTorch 모델 가중치의 기본 형식으로 사용됩니다.

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

pickle란 무엇인가요?

공식 문서에서:

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

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

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

핵심적으로, pickle은 기본적으로 명령어나 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의 역직렬화(de-serialization)는 코드가 실행될 수 있음을 의미합니다. 하지만 여기에는 몇 가지 제한 사항이 있습니다. 최상위 모듈의 함수와 클래스만 참조할 수 있으며, 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 모듈을 가져올 때 임의의 코드가 실행될 수 있습니다
eval 또는 exec과 같은 내장 함수를 가져와 임의의 코드를 실행할 수 있습니다
__init__ 메서드는 객체를 인스턴스화할 때 생성자로 호출될 수 있습니다.

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

완화 전략

pickle을 사용하지 마세요

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

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

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

Hub에서는 GPG 키로 커밋에 서명할 수 있습니다. 이는 파일이 안전하다는 것을 보장하지는 않지만, 파일의 출처는 보장합니다.

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

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

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

ClamAV 스캔
Pickle Import 스캔

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

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

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

이 데이터는 pickletools.genops 덕분에 얻을 수 있으며, 이 함수를 사용하면 잠재적으로 위험한 코드를 실행하지 않고 파일을 읽을 수 있습니다.

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

면책 조항: 이 방법은 100% 완벽하지 않습니다. 사용자가 특정 항목이 안전한지 여부를 확인하는 것이 책임입니다. 우리는 Python 패키지의 안전성을 능동적으로 감사하지 않으며, 안전/불안전한 import 목록은 최선의 노력으로 유지 관리됩니다. 안전하지 않다고 생각하는 항목이 있으면 website at huggingface.co로 이메일을 보내 알려 주시면 해당 항목을 불안전한 것으로 표시하겠습니다.

가능한 해결 방법

Unpickler를 이것처럼 커스텀으로 만들 수 있다고 생각할 수 있습니다. 하지만 이 고도화된 익스플로잇에서 볼 수 있듯이, 이는 작동하지 않습니다.

다행히 eval import의 흔적이 항상 남아 있으므로, opcodes를 직접 읽으면 악성 사용을 감지할 수 있습니다.

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

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

예를 들어 numpy.*와 같은 간단한 한 줄로 모든 numpy 서브모듈을 허용할 수 있는 정규식 유사 형식을 상상해 볼 수 있습니다.

추가 자료

pickle - Python 객체 직렬화 - Python 3.10.6 문서

위험한 Pickles - 악성 Python 직렬화

GitHub - trailofbits/fickling: Python pickling 디컴파일러 및 정적 분석기

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

Python의 pickle 모듈 재활성화

Update on GitHub

←악성 코드 스캐닝 Secrets Scanning→