cntk-poc / vulnerable_code.py

Upload vulnerable_code.py with huggingface_hub

d3735ec verified 5 months ago

2.47 kB

# 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

	# 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