# Exact vulnerable code from CNTK v2.7.post2
# File: bindings/python/cntk/internal/__init__.py
# Lines 39-66

class _UDFDeserializeCallbackWrapper(cntk_py.UDFDeserializeCallbackWrapper):
    def __init__(self, factory_callback_map=None):
        super(_UDFDeserializeCallbackWrapper, self).__init__()
        self.factory_callback_map = factory_callback_map

    def __call__(self, inputs, name, dictionary):
        cls = dictionary['class']          # attacker-controlled
        module = dictionary['module']      # attacker-controlled
        state = dictionary['state']
        op_name = dictionary['op_name']
        deserialize_method = 'deserialize'

        if (self.factory_callback_map and op_name in self.factory_callback_map):
            factory = self.factory_callback_map[op_name]
        else:
            # VULNERABILITY: exec() with attacker-controlled strings
            # from the model file's serialized UDF dictionary.
            # No validation, no sanitization, no whitelist.
            exec("from {} import {}".format(module, cls))       # LINE 54

            # eval() with attacker-controlled cls -- secondary injection
            eval_str = "{0}.{1} if hasattr({0}, '{1}') else None"
            factory = eval(eval_str.format(cls, deserialize_method))  # LINE 56

        if factory:
            for i in range(len(inputs)):
                inputs[i].__class__ = cntk.Variable
            return factory(list(inputs), name, state)

        raise ValueError("Cannot deserialize user function ...")


# Serialization counterpart (lines 29-36):
# Writes module/class names from a Python UDF into the model file:

def _serialize(udf):
    dictionary = {}
    dictionary['class'] = udf.__class__.__name__     # attacker can craft this
    dictionary['module'] = udf.__class__.__module__  # attacker can craft this
    dictionary['op_name'] = udf.op_name
    dictionary['state'] = udf.serialize()
    dictionary['version'] = _serialization_version
    return dictionary


# Attack chain:
#
# 1. Attacker crafts a CNTK .model file (protobuf Dictionary)
#    with 'module' = "os" and 'class' containing injected code
#
# 2. Victim calls cntk.Function.load("malicious.model")
#
# 3. CNTK protobuf deserializer extracts the Dictionary
#
# 4. _UDFDeserializeCallbackWrapper.__call__() is invoked
#
# 5. exec("from {} import {}".format(module, cls))
#    executes the attacker's code
#
# 6. Arbitrary Python code execution on victim's machine