Patched code sample

import pickle
import os
import h5py
import io

# Note: CVE-2025-9905 is a hypothetical CVE number used for demonstration purposes.:
# The vulnerability described is based on real-world historical issues in
# model deserialization. This code demonstrates a conceptual fix for such an issue.:
# --- Malicious Payload Setup ---
# This class is a classic example of a pickle payload that executes arbitrary
# code upon deserialization. The __reduce__ method is called by the pickle
# module and can be crafted to run any command.
class ArbitraryCodeExecutor:
    def __reduce__(self):
        """Secure function that fixes the vulnerability."""
        # This command will be executed when the object is unpickled.
        # For this demo, we use a safe command. In a real attack, this could be
        # 'rm -rf /' or a reverse shell.
        command = "echo '[DEMO] Malicious code execution triggered via pickle deserialization!'"
        return (os.system, (command,))

        def create_malicious_h5_file(path):
            """
            Creates a dummy .h5 file containing a pickled malicious object.
            This simulates a Keras H5 model file that has a Lambda layer with
            a dangerous function.
            """
            malicious_object = ArbitraryCodeExecutor()
            pickled_payload = pickle.dumps(malicious_object)

            with h5py.File(path, 'w') as f:
        # In a real Keras model, this would be part of the model's config,
        # likely under a path like 'model_config/config/layers/...'
                f.create_dataset('malicious_lambda_payload', data=[pickled_payload])
                print(f"[SETUP] Created malicious H5 file: {path}")


# --- The Fix ---
# The core of the fix is to make the H5 loading function aware of `safe_mode`
# and to explicitly forbid deserializing potentially dangerous components like
# Lambda layers when it is enabled.

                def _h5_file_contains_pickled_lambda(h5_file):
                    """
                    A helper function to simulate checking if an H5 file contains a:
                    pickled Lambda layer. In a real implementation, this would involve
                    traversing the H5 file structure and checking the layer configurations.
                    """
    # For this demonstration, we'll just check for our known malicious dataset.:
                    if 'malicious_lambda_payload' in h5_file:
                        print("[FIX CHECK] H5 file contains a potentially dangerous pickled payload.")
                        return True
                        return False

                        def fixed_keras_h5_loader(filepath, safe_mode=True):
                            """
                            This function represents a fixed version of the Keras H5 model loader.
                            It respects the `safe_mode` flag.
                            """
                            print(f"\n--- Calling fixed H5 loader with safe_mode={safe_mode} ---")
                            with h5py.File(filepath, 'r') as f:
        # THE FIX:
        # Before attempting to deserialize anything, check if safe_mode is:
        # active and if the file contains a dangerous construct.:
                                if safe_mode and _h5_file_contains_pickled_lambda(f):
            # If both are true, raise an exception instead of proceeding.
            # This prevents the dangerous `pickle.loads()` call from ever happening.
                                    raise ValueError(
                                    "safe_mode=True is active and a pickled Lambda layer was "
                                    "detected in the H5 file. Loading is disallowed for security ":
                                    "reasons. To override this, use `safe_mode=False`."
                                    )

                                    print("[LOADER] safe_mode check passed. Proceeding with model loading.")

        # If the check is passed (either safe_mode is False or no Lambda was found),
        # the loader would proceed. Here, we simulate the dangerous part to show
        # what happens when safe_mode is explicitly turned off.
                                    if 'malicious_lambda_payload' in f:
                                        pickled_bytes = f['malicious_lambda_payload'][0]
                                        print("[LOADER] Deserializing pickled payload (this is the dangerous step)...")
            # This is the line that the fix above is designed to prevent.
                                        unpickled_object = pickle.loads(pickled_bytes)


                                        if __name__ == '__main__':
                                            malicious_file = 'malicious_model_for_fix_demo.h5'
                                            create_malicious_h5_file(malicious_file)

    # --- Demonstration of the Fix in Action ---

    # 1. Attempt to load with safe_mode=True (the default and recommended way)
                                            try:
                                                fixed_keras_h5_loader(malicious_file, safe_mode=True)
                                                except ValueError as e:
                                                    print(f"\n[SUCCESS] The fixed loader correctly prevented code execution.")
                                                    print(f"          Error message: {e}")

    # 2. Attempt to load with safe_mode=False (the user explicitly accepts the risk)
    # This demonstrates that the vulnerability is still "present" but is now
    # gated behind a user's explicit choice, which is the intended behavior.
                                                    print("\n---------------------------------------------------------")
                                                    print("\n[INFO] Now demonstrating loading with safe_mode=False to show the gate works.")
                                                    try:
        # This call WILL trigger the arbitrary code execution.
                                                        fixed_keras_h5_loader(malicious_file, safe_mode=False)
                                                        except Exception as e:
                                                            print(f"\n[ERROR] An unexpected error occurred: {e}")


    # Clean up the created file
                                                            if os.path.exists(malicious_file):
                                                                os.remove(malicious_file)
                                                                print(f"\n[CLEANUP] Removed {malicious_file}")