Analyse des Pickles

Pickle est un format de sérialisation largement utilisé en ML. En particulier, c'est le format par défaut pour les poids des modèles PyTorch.

Il existe des attaques par exécution de code arbitraire dangereuses qui peuvent être menées lors du chargement d'un fichier pickle. Nous recommandons de charger des modèles provenant d'utilisateurs et d'organisations de confiance, de vous fier à des commits signés et/ou de charger des modèles au format TF ou Jax grâce au mécanisme de conversion automatique from_tf=True. Nous atténuons également ce problème en affichant/« vérifiant » la liste des imports de tout fichier picklé, directement sur le Hub. Enfin, nous expérimentons un nouveau format de sérialisation simple pour les poids, appelé safetensors.

Qu'est-ce qu'un pickle ?

D'après la documentation officielle :

Le module pickle implémente des protocoles binaires pour la sérialisation et la désérialisation d'une structure d'objets Python.

Cela signifie que pickle est un protocole de sérialisation, un outil que l'on utilise pour partager efficacement des données entre différentes parties.

Nous appelons un pickle le fichier binaire généré lors de la sérialisation.

Au cœur du processus, le pickle est essentiellement une pile d’instructions ou d’opcodes. Comme vous l’avez probablement deviné, il n’est pas lisible par un humain. Les opcodes sont générés lors de la sérialisation (pickling) et lus séquentiellement lors de la désérialisation (unpickling). En fonction de l’opcode, une action donnée est exécutée.

Voici un petit exemple :

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)

Lorsque vous exécutez cela, un fichier pickle est créé et les instructions suivantes sont affichées dans votre terminal :

    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

Ne vous inquiétez pas trop des instructions pour l'instant, sachez simplement que le module pickletools est très utile pour analyser les pickles. Il vous permet de lire les instructions du fichier sans exécuter de code.

Pickle n’est pas simplement un protocole de sérialisation, il offre plus de flexibilité en permettant aux utilisateurs d’exécuter du code Python lors de la désérialisation. Cela ne semble pas très rassurant, n’est-ce pas ?

Pourquoi est-ce dangereux ?

Comme nous l’avons indiqué plus haut, la désérialisation de pickle signifie que du code peut être exécuté. Mais cela comporte certaines limitations : vous ne pouvez référencer que des fonctions et des classes du module de niveau supérieur ; vous ne pouvez pas les intégrer dans le fichier pickle lui-même.

Retour à la case départ :

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)

Lorsque nous exécutons ce script, nous obtenons à nouveau payload.pkl. Lorsque nous vérifions le contenu du fichier :

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

On peut y voir qu’il n’y a pas grand-chose, quelques opcodes et les données associées. Vous vous demandez peut-être quel est le problème avec pickle ?

Essayons autre chose :

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)

Ici, nous utilisons la bibliothèque fickling pour la simplicité. Elle nous permet d’ajouter des instructions pickle pour exécuter du code contenu dans une chaîne de caractères via la fonction exec. C’est ainsi que vous contournez le fait que vous ne pouvez pas définir de fonctions ou de classes dans vos pickles : vous exécutez exec sur du code Python enregistré sous forme de chaîne.

Lors de l'exécution de ce code, un fichier payload.pkl est créé et le texte suivant est affiché :

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

Si nous vérifions le contenu du fichier pickle, nous obtenons :

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

Concrètement, voici ce qui se passe lors du dépickling :

# ...
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))
# ...

Les instructions qui représentent une menace sont STACK_GLOBAL, GLOBAL et 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.

Pour résumer, pickle est dangereux car :

• lors de l'importation d'un module Python, du code arbitraire peut être exécuté
• vous pouvez importer des fonctions intégrées comme eval ou exec, qui peuvent être utilisées pour exécuter du code arbitraire
• lors de l'instantiation d'un objet, le constructeur peut être appelé

C’est pourquoi la plupart des documents utilisant pickle indiquent de ne pas désérialiser de données provenant de sources non fiables.

Stratégies d'atténuation

Ne pas utiliser pickle

Conseil judicieux, Luc, mais pickle est utilisé abondamment et ne disparaîtra pas de sitôt : trouver un nouveau format qui convienne à tous et initier le changement prendra du temps.

Que pouvons-nous faire pour l'instant ?

Charger des fichiers à partir d'utilisateurs et d'organisations de confiance

Sur le Hub, vous avez la possibilité de signer vos commits avec une clé GPG. Cela ne garantit pas que votre fichier est sûr, mais cela garantit l'origine du fichier.

Si vous connaissez et faites confiance à l’utilisateur A et que l’engagement contenant le fichier sur le Hub est signé par la clé GPG de l’utilisateur A, il est assez sûr de supposer que vous pouvez faire confiance au fichier.

Charger les poids du modèle à partir de TF ou Flax

Les points de contrôle TensorFlow et Flax ne sont pas affectés et peuvent être chargés dans des architectures PyTorch à l’aide des arguments from_tf et from_flax de la méthode from_pretrained pour contourner ce problème.

E.g.:

from transformers import AutoModel

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

Utilisez votre propre format de sérialisation

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

Ce dernier format, safetensors, est un format de sérialisation simple sur lequel nous travaillons et avec lequel nous expérimentons actuellement ! N’hésitez pas à nous aider ou à contribuer si vous le pouvez 🔥.

Améliorer torch.load/save

Une discussion est en cours sur PyTorch concernant la mise en place d’une manière sûre de charger uniquement les poids depuis un fichier *.pt par défaut – n’hésitez pas à y participer !

Le scanner de sécurité du Hub

Ce que nous avons maintenant

Nous avons créé un analyseur de sécurité qui examine chaque fichier poussé sur le Hub et exécute des vérifications de sécurité. À l'heure où nous écrivons ces lignes, il effectue deux types d'analyses :

• Scans ClamAV
• Analyses des imports Pickle

Pour les analyses ClamAV, les fichiers sont passés au crible de l’antivirus open source ClamAV. Bien que cela couvre une large part des fichiers dangereux, cela ne protège pas contre les exploits pickle.

Nous avons mis en place une analyse d’importation Pickle, qui extrait la liste des imports référencés dans un fichier pickle. Chaque fois que vous téléversez un pytorch_model.bin ou tout autre fichier picklé, cette analyse est exécutée.

Sur le hub, la liste des imports est affichée à côté de chaque fichier contenant des imports. Si un import semble suspect, il est mis en évidence.

Nous obtenons ces données grâce à pickletools.genops, qui nous permet de lire le fichier sans exécuter de code potentiellement dangereux.

Notez que c'est ce qui permet de savoir si, lors du dépickling d'un fichier, il va REDUCE sur une fonction potentiellement dangereuse qui a été importée par *GLOBAL.

Avertissement : ceci n’est pas à 100 % infaillible. Il est de votre responsabilité, en tant qu’utilisateur, de vérifier si quelque chose est sûr ou non. Nous n’auditionnons pas activement les paquets Python pour la sécurité ; les listes d’imports sûrs/insécurisés dont nous disposons sont maintenues sur une base de meilleur effort. Veuillez nous contacter si vous pensez que quelque chose n’est pas sûr, et nous le signalerons comme tel, en nous envoyant un e-mail à website at huggingface.co

Solutions potentielles

On pourrait envisager de créer un Unpickler personnalisé, à l’image de celui-ci. Mais comme on peut le voir dans cette exploitation sophistiquée, cela ne fonctionnera pas.

Heureusement, il y a toujours une trace de l'import eval, de sorte que la lecture directe des opcodes devrait permettre de détecter un usage malveillant.

La solution que je propose actuellement consiste à créer un fichier ressemblant à un .gitignore, mais destiné aux imports.

Ce fichier serait une liste blanche d’imports qui marquerait un fichier pytorch_model.bin comme dangereux s’il contient des imports non inclus dans la liste blanche.

On pourrait imaginer un format de type regex permettant, par exemple, d’autoriser tous les sous-modules de numpy via une simple ligne comme : numpy.*.

Lecture complémentaire

pickle - Sérialisation d'objets Python - Documentation Python 3.10.6

Pickles dangereux - Sérialisation Python malveillante

GitHub - trailofbits/fickling : un décompilateur et un analyseur statique pour le pickling Python

Exploitation des pickles Python

cpython/pickletools.py at 3.10 · python/cpython

cpython/pickle.py at 3.10 · python/cpython

CrypTen/serial.py sur main · facebookresearch/CrypTen

CTFtime.org / Balsn CTF 2019 / pyshv1 / Writeup

Réhabiliter le module pickle de Python