Patched code sample

import argparse
from transformers import AutoModelForCausalLM

def load_model_securely():
    """
    Demonstrates the fix for a vulnerability where `trust_remote_code` was
    hardcoded to True.

    The fix involves making `trust_remote_code` a user-configurable command-line
    argument that defaults to False, ensuring users must explicitly opt-in
    to the dangerous behavior.
    """
    parser = argparse.ArgumentParser(
        description="Securely load a model from HuggingFace Hub."
    )
    parser.add_argument(
        "--model-name",
        type=str,
        required=True,
        help="The repository ID of the model on HuggingFace Hub (e.g., 'gpt2').",
    )

    # THE FIX: Add a command-line flag that defaults to False.
    # The vulnerability existed because the value was previously hardcoded to True.
    parser.add_argument(
        "--trust-remote-code",
        action="store_true",  # Makes the default value False
        help="If provided, trust remote code found in the model. (UNSAFE)",
    )

    args = parser.parse_args()

    print(f"Attempting to load model: {args.model_name}")
    print(f"Configuration: trust_remote_code={args.trust_remote_code}")

    try:
        # VULNERABLE (old) code would have looked like this:
        # model = AutoModelForCausalLM.from_pretrained(
        #     args.model_name,
        #     trust_remote_code=True  # <-- Hardcoded True is the vulnerability
        # )

        # FIXED code uses the user-provided flag, which is False by default.
        model = AutoModelForCausalLM.from_pretrained(
            args.model_name,
            trust_remote_code=args.trust_remote_code,
        )
        print("Model loading function executed successfully.")

    except Exception as e:
        print(f"Model loading failed. This is expected if the model is malicious or does not exist.")
        print(f"Error: {e}")


if __name__ == "__main__":
    # To run safely (default behavior):
    # python <filename>.py --model-name="gpt2"
    #
    # To run with the potentially insecure option (requires explicit user action):
    # python <filename>.py --model-name="some/malicious-model" --trust-remote-code
    load_model_securely()