Сканирование Pickle

Pickle — это широко используемый формат сериализации в машинном обучении. В частности, это формат по умолчанию для весов моделей PyTorch.

Существуют опасные атаки с произвольным выполнением кода, которые могут быть реализованы при загрузке pickle-файла. Мы рекомендуем загружать модели только от пользователей и организаций, которым вы доверяете, полагаться на подписанные коммиты и/или загружать модели из форматов TF или Jax с использованием механизма автоматического преобразования from_tf=True. Кроме того, мы смягчаем эту проблему, отображая и «проверяя» список импортов в любом pickle-файле непосредственно на Hub. Наконец, мы экспериментируем с новым простым форматом сериализации весов, названным safetensors.

Что такое pickle?

Из официальной документации :

Модуль pickle реализует бинарные протоколы для сериализации и десериализации структуры объектов Python.

Это означает, что pickle представляет собой протокол сериализации, который используется для эффективного обмена данными между участниками.

Файл, сгенерированный в процессе pickling, мы называем pickle.

В основе pickle лежит, по сути, стек инструкций или операционных кодов. Как вы, вероятно, догадались, он не предназначен для чтения человеком. Операционные коды генерируются при сериализации и считываются последовательно при десериализации. В зависимости от операционного кода выполняется соответствующее действие.

Вот небольшой пример:

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

Мы видим, что там не так много: несколько opcodes и связанные с ними данные. Возможно, вы думаете, в чём же проблема с 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. Она позволяет добавлять в pickle инструкции для выполнения кода, содержащегося в строке, с помощью функции exec. Таким образом вы обходите ограничение на определение функций или классов внутри pickle: вы запускаете exec на Python-коде, сохранённом в виде строки.

При запуске этого кода создаётся файл 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 is what tells the unpickler to execute the function with the provided arguments and *GLOBAL instructions are telling the unpickler to import stuff.

Подводя итог, pickle опасен, потому что:

• при импорте модуля Python может быть выполнен произвольный код
• вы можете импортировать встроенные функции, такие как eval или exec, которые можно использовать для выполнения произвольного кода
• при создании объекта может быть вызван конструктор

Именно поэтому в большинстве документов, использующих pickle, указано, что не следует выполнять распаковку данных из ненадёжных источников.

Стратегии смягчения

Не используйте pickle

Совет, конечно, разумный, Luc, но pickle используется повсеместно и никуда не денется в ближайшее время: чтобы найти новый формат, который устроит всех, и инициировать переход, потребуется время.

Что же мы можем сделать на данный момент?

Загрузка файлов от пользователей и организаций, которым вы доверяете

На Hub вы можете подписывать свои коммиты с помощью GPG-ключа. Это не гарантирует, что ваш файл безопасен, но гарантирует его происхождение.

Если вы знаете и доверяете пользователю A, а коммит, содержащий файл на Hub, подписан GPG-ключом пользователя A, можно с высокой степенью уверенности считать, что файлу можно доверять.

Загрузка весов модели из TF или Flax

Чекпоинты TensorFlow и Flax не затрагиваются этой проблемой и могут быть загружены в архитектуры PyTorch с помощью аргументов from_tf и from_flax для метода from_pretrained.

E.g.:

from transformers import AutoModel

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

Используйте собственный формат сериализации

• MsgPack
• Protobuf
• Cap’n’proto
• Avro
• safetensors

Этот последний формат, safetensors, — это простой формат сериализации, над которым мы сейчас работаем и который тестируем! Пожалуйста, помогите или вносите свой вклад, если можете 🔥.

Улучшение torch.load/save

В PyTorch сейчас ведётся открытое обсуждение о том, чтобы по умолчанию безопасно загружать только веса из файлов *.pt – присоединяйтесь к дискуссии!

Сканер безопасности Hub

Что у нас есть сейчас

Мы создали сканер безопасности, который проверяет каждый файл, загруженный в Hub, и выполняет проверки безопасности. На момент написания документации он выполняет два типа сканирования:

• Сканирование ClamAV
• Сканирование импорта Pickle

Для сканирования ClamAV файлы обрабатываются с помощью открытого антивируса ClamAV. Хотя это покрывает значительное количество опасных файлов, оно не защищает от эксплойтов pickle.

Мы реализовали сканирование Pickle Import, которое извлекает список импортов, упомянутых в pickle-файле. Каждый раз, когда вы загружаете pytorch_model.bin или любой другой pickle-файл, выполняется это сканирование.

На хабе список импортов отображается рядом с каждым файлом, содержащим импорты. Если какой-либо импорт выглядит подозрительно, он будет выделен.

Мы получаем эти данные благодаря pickletools.genops, что позволяет читать файл без выполнения потенциально опасного кода.

Обратите внимание, что именно это позволяет определить, будет ли при распаковке файла вызвана функция REDUCE для потенциально опасной функции, импортированной с помощью *GLOBAL.

Дисклеймер: это не на 100% надёжный метод. Ответственность за проверку безопасности того или иного кода лежит на пользователе. Мы не проводим активный аудит пакетов Python на предмет безопасности; списки безопасных и небезопасных импортов поддерживаются на основе лучших усилий. Если вы считаете, что что-то небезопасно, а мы отмечаем это как безопасное, свяжитесь с нами, отправив письмо на website at huggingface.co

Возможные решения

Можно было бы подумать о создании собственного Unpickler, похожего на этот. Но, как мы видим в этом сложном эксплойте, это не сработает.

К счастью, всегда остаётся след импорта eval, поэтому прямое чтение операндов должно позволять выявлять вредоносное использование.

Решение, которое я предлагаю на данный момент, — это создание файла, похожего на .gitignore, но для импортов.

Этот файл представляет собой белый список импортов, при наличии импортов, не входящих в этот список, файл pytorch_model.bin будет помечен как опасный.

Можно представить себе формат, похожий на регулярные выражения, в котором, например, можно разрешить все подмодули 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 в main · facebookresearch/CrypTen

CTFtime.org / Balsn CTF 2019 / pyshv1 / Разбор

Реабилитация модуля pickle в Python