Patched code sample

import os
import yaml
from pathlib import Path

# The CVE describes the vulnerability in a file like this, located at:
# `feast/sdk/python/feast/infra/compute_engines/kubernetes/main.py`
#
# For demonstration, we assume these paths are mounted into the pod.
FEAST_CONFIG_DIR = Path("/tmp/feast_demo/")  # Using /tmp for a runnable example
FEATURE_STORE_YAML_PATH = FEAST_CONFIG_DIR / "feature_store.yaml"
MATERIALIZATION_CONFIG_YAML_PATH = FEAST_CONFIG_DIR / "materialization_config.yaml"


def _load_and_apply_configuration():
    """
    Loads configuration from YAML files. This function demonstrates the fix
    for CVE-2025-11157 by using a safe YAML loader.

    The vulnerability existed because the original code used `yaml.load(..., Loader=yaml.Loader)`,
    which allows for arbitrary Python object instantiation and code execution.
    """
    # VULNERABLE CODE (for reference, this is what was fixed):
    #
    # with open(FEATURE_STORE_YAML_PATH, "r") as f:
    #     # Unsafe deserialization allows an attacker to run OS commands.
    #     feature_store_config = yaml.load(f, Loader=yaml.Loader)
    #
    # with open(MATERIALIZATION_CONFIG_YAML_PATH, "r") as f:
    #     # This was also vulnerable.
    #     materialization_config = yaml.load(f, Loader=yaml.Loader)

    # --- FIXED CODE ---
    # The fix is to replace the unsafe `yaml.load` with `yaml.safe_load`.
    # `yaml.safe_load()` is equivalent to `yaml.load(..., Loader=yaml.SafeLoader)`.
    # The `SafeLoader` only parses standard YAML tags and prevents the
    # instantiation of arbitrary Python objects, mitigating the RCE vulnerability.
    try:
        print(f"Safely loading configuration from {FEATURE_STORE_YAML_PATH}...")
        with open(FEATURE_STORE_YAML_PATH, "r") as f:
            feature_store_config = yaml.safe_load(f)

        print(f"Safely loading configuration from {MATERIALIZATION_CONFIG_YAML_PATH}...")
        with open(MATERIALIZATION_CONFIG_YAML_PATH, "r") as f:
            materialization_config = yaml.safe_load(f)

        print("\nSuccessfully and safely loaded configuration.")
        # The `safe_load` function would have raised an error if it encountered a
        # malicious tag like `!!python/object/apply`, preventing code execution.

    except yaml.YAMLError as e:
        print(f"\n[SUCCESS] The fix worked! Caught a potentially malicious YAML tag.")
        print(f"Error details: {e}")
        print("The unsafe object was NOT instantiated, and no command was executed.")


def main():
    """
    Main entry point to simulate the environment and demonstrate the fix.
    """
    print("--- Demonstrating fix for CVE-2025-11157 ---")

    # 1. Simulate the environment by creating the config directory.
    os.makedirs(FEAST_CONFIG_DIR, exist_ok=True)

    # 2. Create a legitimate feature_store.yaml file.
    legit_fs_yaml = """
    project: my_feast_project
    registry: data/registry.db
    provider: gcp
    """
    with open(FEATURE_STORE_YAML_PATH, "w") as f:
        f.write(legit_fs_yaml)

    # 3. Create a malicious materialization_config.yaml.
    # An attacker with write access to this file could inject this payload.
    # The vulnerable code would execute `os.system("touch /tmp/pwned")`.
    malicious_yaml_payload = """
    !!python/object/apply:os.system
    - "touch /tmp/pwned_by_cve"
    """
    with open(MATERIALIZATION_CONFIG_YAML_PATH, "w") as f:
        f.write(malicious_yaml_payload)

    print(f"Created a malicious payload in {MATERIALIZATION_CONFIG_YAML_PATH}")
    print("Payload tries to execute: os.system('touch /tmp/pwned_by_cve')\n")

    # 4. Run the fixed configuration loader.
    _load_and_apply_configuration()

    # 5. Verify that the malicious command was NOT executed.
    pwned_file = Path("/tmp/pwned_by_cve")
    if not pwned_file.exists():
        print("\n[VERIFICATION PASSED] The file /tmp/pwned_by_cve was not created.")
        print("The remote code execution was successfully prevented.")
    else:
        print("\n[VERIFICATION FAILED] The file /tmp/pwned_by_cve was created.")
        os.remove(pwned_file)

    # --- Cleanup ---
    os.remove(FEATURE_STORE_YAML_PATH)
    os.remove(MATERIALIZATION_CONFIG_YAML_PATH)
    os.rmdir(FEAST_CONFIG_DIR)


if __name__ == "__main__":
    main()