Hub documentation
Escaneo de Pickle
Escaneo de Pickle
Pickle es un formato de serialización ampliamente utilizado en ML. En particular, es el formato predeterminado para los pesos de los modelos de PyTorch.
Existen ataques de ejecución de código arbitrario y peligroso que pueden perpetrarse al cargar un archivo pickle. Sugerimos cargar modelos de usuarios y organizaciones de confianza, confiar en commits firmados y/o cargar modelos desde formatos de TF o Jax con el mecanismo de autoconversión from_tf=True. También mitigamos este problema mostrando/“revisando” la lista de imports en cualquier archivo pickle, directamente en el Hub. Por último, estamos experimentando con un nuevo formato de serialización simple para pesos llamado safetensors.
¿Qué es un pickle?
De la documentación oficial :
El módulo
pickleimplementa protocolos binarios para serializar y deserializar una estructura de objetos de Python.
Esto significa que pickle es un protocolo de serialización, algo que se utiliza para compartir datos de manera eficiente entre distintas partes.
Llamamos pickle al archivo binario que se generó durante el proceso de pickling.
En su núcleo, el pickle es básicamente una pila de instrucciones u opcodes. Como probablemente haya adivinado, no es legible para humanos. Los opcodes se generan al hacer pickle y se leen secuencialmente al deshacer el pickle. Según el opcode, se ejecuta una acción dada.
Aquí hay un pequeño ejemplo:
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)Cuando ejecutes esto, se creará un archivo pickle y se imprimirán las siguientes instrucciones en tu 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 = 4No te preocupes demasiado por las instrucciones por ahora, solo ten en cuenta que el módulo pickletools es muy útil para analizar pickles. Te permite leer las instrucciones del archivo sin ejecutar ningún código.
Pickle no es simplemente un protocolo de serialización, ya que ofrece mayor flexibilidad al permitir a los usuarios ejecutar código de Python en el momento de la deserialización. ¿No suena bien, verdad?
¿Por qué es peligroso?
Como hemos indicado anteriormente, deserializar pickle significa que se puede ejecutar código. Pero esto conlleva ciertas limitaciones: solo se pueden referenciar funciones y clases del módulo de nivel superior; no se pueden incrustar en el propio archivo pickle.
De vuelta a la tabula rasa:
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)Cuando ejecutamos este script, obtenemos de nuevo payload.pkl. Al revisar el contenido del archivo:
# 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. │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘Podemos ver que no hay mucho contenido, solo algunas opcodes y los datos asociados. Es posible que estés pensando: ¿cuál es el problema con pickle?
Probemos algo más:
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)Aquí utilizamos la biblioteca fickling por simplicidad. Nos permite añadir instrucciones de pickle para ejecutar código contenido en una cadena mediante la función exec. Así es como se evita el hecho de que no se pueden definir funciones o clases en los pickles: se ejecuta exec sobre código de python guardado como una cadena.
Al ejecutar esto, se crea un payload.pkl y se imprime lo siguiente:
you've been pwned !
my friend needs to know thisSi comprobamos el contenido del archivo pickle, obtenemos:
# 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×. │
└────────┴─────────────────────────┴─────────────────────────┴────────┴────────┘Básicamente, esto es lo que ocurre cuando se deserializa:
# ...
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))
# ...Las instrucciones que suponen una amenaza son STACK_GLOBAL, GLOBAL y REDUCE.
REDUCE es lo que le indica al deserializador que ejecute la función con los argumentos proporcionados, y las instrucciones *GLOBAL le indican al deserializador que import elementos.
Para resumir, pickle es peligroso porque:
- al importar un módulo de Python, se puede ejecutar código arbitrario
- puede importar funciones integradas como
evaloexec, que se pueden usar para ejecutar código arbitrario - al instanciar un objeto, se puede llamar al constructor
Por eso se indica en la mayoría de la documentación que utiliza pickle: no deserialice datos de fuentes no confiables.
Estrategias de mitigación
No uses pickle
Consejo sensato, Luc, pero pickle se usa profusamente y no va a desaparecer pronto: encontrar un nuevo formato que satisfaga a todos e iniciar el cambio llevará algún tiempo.
¿Qué podemos hacer por ahora?
Carga archivos de usuarios y organizaciones en los que confíes
En el Hub, puedes firmar tus commits con una clave GPG. Esto no garantiza que tu archivo sea seguro, pero sí garantiza el origen del archivo.
Si conoces y confías en el usuario A y el commit que incluye el archivo en el Hub está firmado con la clave GPG del usuario A, es bastante seguro asumir que puedes confiar en el archivo.
Cargar los pesos del modelo desde TF o Flax
Los puntos de control de TensorFlow y Flax no se ven afectados y pueden cargarse en arquitecturas de PyTorch utilizando los argumentos from_tf y from_flax del método from_pretrained para evitar este problema.
E.g.:
from transformers import AutoModel
model = AutoModel.from_pretrained("google-bert/bert-base-cased", from_flax=True)Usa tu propio formato de serialización
Este último formato, safetensors, es un formato de serialización simple en el que estamos trabajando y experimentando actualmente. ¡Por favor, ayuda o contribuye si puedes! 🔥
Mejorar torch.load/save
Hay una discusión abierta en curso en PyTorch sobre tener una forma segura de cargar solo los pesos desde un archivo *.pt de forma predeterminada ¡participa allí!
Hub’s Security Scanner
Lo que tenemos ahora
Hemos creado un escáner de seguridad que analiza cada archivo enviado al Hub y ejecuta comprobaciones de seguridad. En el momento de redactar este texto, ejecuta dos tipos de escaneos:
- Escaneos de ClamAV
- Escaneos de importación de Pickle
En las escaneos de ClamAV, los archivos se procesan con el antivirus de código abierto ClamAV. Si bien esto cubre una gran cantidad de archivos peligrosos, no cubre las explotaciones de pickle.
Hemos implementado un análisis de importación Pickle, que extrae la lista de importaciones referenciadas en un archivo pickle. Cada vez que subes un pytorch_model.bin o cualquier otro archivo pickle, se ejecuta este análisis.
En el hub, la lista de importaciones se mostrará junto a cada archivo que contenga importaciones. Si alguna importación parece sospechosa, se resaltará.
Obtenemos estos datos gracias a pickletools.genops, que nos permite leer el archivo sin ejecutar código potencialmente peligroso.
Tenga en cuenta que esto es lo que permite saber si, al deserializar un archivo, se ejecutará REDUCE sobre una función potencialmente peligrosa que fue importada mediante *GLOBAL.
Aviso legal: esto no es 100% infalible. Es responsabilidad del usuario verificar si algo es seguro o no. No auditamos activamente los paquetes de Python en busca de seguridad; las listas de importaciones seguras/insaludables que tenemos se mantienen de la mejor manera posible. Si crees que algo no es seguro, contáctanos para que lo marquemos como tal, enviándonos un correo electrónico a website at huggingface.co
Soluciones potenciales
Podría pensarse en crear un Unpickler personalizado, similar a este. Pero, como podemos ver en este explotado sofisticado, esto no funcionará.
Afortunadamente, siempre queda un rastro de la importación de eval, por lo que leer los opcodes directamente debería permitir detectar un uso malicioso.
La solución que propongo actualmente es crear un archivo similar a un .gitignore pero para las importaciones.
Este archivo sería una lista blanca de importaciones que marcaría un archivo pytorch_model.bin como peligroso si hay importaciones no incluidas en la lista blanca.
Podría imaginarse un formato similar a una expresión regular que, por ejemplo, permitiera todos los submódulos de numpy mediante una simple línea como: numpy.*.
Lecturas adicionales
pickle - Serialización de objetos en Python - Documentación de Python 3.10.6
Pickles peligrosos: serialización maliciosa de Python
GitHub - trailofbits/fickling: Un descompilador y analizador estático de pickling en Python
cpython/pickletools.py en 3.10 · python/cpython
cpython/pickle.py en 3.10 · python/cpython
CrypTen/serial.py en main · facebookresearch/CrypTen
CTFtime.org / Balsn CTF 2019 / pyshv1 / Writeup
Rehabilitar el módulo pickle de Python
Update on GitHub