Vulnerable code sample

# This code provides a simplified representation of the vulnerable logic
# found in the `diffusers` library prior to version 0.38.0.

def _find_and_read_remote_file(repo_id, filename):
    """
    Simulates finding a malicious file in a remote repository.
    An attacker would host a repository containing a file named "None.py".
    """
    if filename == "None.py":
        # This string represents the content of the attacker's malicious file.
        return """
print("!!! VULNERABILITY TRIGGERED: Remote Code Execution Successful !!!")
print("This code is running from the attacker's 'None.py' file.")

# The file must contain a class that subclasses the pipeline.
class MaliciousPipeline:
    def __init__(self, *args, **kwargs):
        print("MaliciousPipeline object has been instantiated.")
"""
    return None

def _resolve_custom_pipeline_and_cls(repo_id, custom_pipeline):
    """
    Simplified representation of the function in pipeline_loading_utils.py.
    """
    # THE VULNERABILITY:
    # When `custom_pipeline` is None, f-string interpolation creates the
    # literal filename "None.py".
    module_file = f"{custom_pipeline}.py"
    
    # Simulate attempting to load the resolved file.
    code_to_execute = _find_and_read_remote_file(repo_id, module_file)
    
    if code_to_execute:
        # In the actual library, `importlib` is used to load the module,
        # which executes the code within it upon import. We use `exec()`
        # here to demonstrate the code execution.
        exec(code_to_execute, {"__name__": "__main__"})


class DiffusionPipeline:
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        # When a user calls this function, `custom_pipeline` is not provided
        # in kwargs, so its value is None.
        custom_pipeline = kwargs.get("custom_pipeline")
        trust_remote_code = kwargs.get("trust_remote_code", False)

        # THE BYPASS: The safety check for `trust_remote_code` is not triggered
        # because the condition `custom_pipeline is not None` evaluates to False.
        if custom_pipeline is not None and not trust_remote_code:
            raise ValueError(
                "Executing custom code requires `trust_remote_code=True`."
            )
        
        # The program proceeds to the vulnerable function call path.
        _resolve_custom_pipeline_and_cls(
            pretrained_model_name_or_path,
            custom_pipeline=custom_pipeline,
        )


# --- DEMONSTRATION OF THE ATTACK ---
# A victim attempts to load what they believe is a standard model.
# They do NOT specify `trust_remote_code=True` and assume they are safe.
if __name__ == "__main__":
    print("Simulating a user calling: DiffusionPipeline.from_pretrained('attacker/repo')\n")
    DiffusionPipeline.from_pretrained("attacker/repo")
    print("\nScript finished.")