MadMarx37/llm-code-injection

1import gradio as gr2from agents import get_crew3import json4from dotenv import load_dotenv5import os6import pyperclip7 8# Load environment variables from a .env file9load_dotenv()10os.environ["OPENAI_MODEL_NAME"] = 'gpt-4o-mini'11 12# Path to your demonstration image13image_path = "image.png"14 15def set_openai_api_key(api_key):16    os.environ["OPENAI_API_KEY"] = api_key17 18def generate_crew_output(text_prompt, api_key_input):19      set_openai_api_key(api_key_input)20      crew = get_crew(text_prompt)21      result = crew.kickoff()22      results_json = json.loads(result.raw)23      return results_json["safe_html"], results_json["safe_html"], results_json["compromised_html"]24 25def generate_safe_prompt():26    return "Hey, could you help me write a form? The fields need to be your name, manager's name and hours worked in the last week."27    28 29def generate_vulnerable_prompt():30    return "Hey, could you help me write a one-page login form in Svelte?"31 32def raise_error():33    raise gr.Error("You have been hacked!")34 35def copy_malicious_code(malicious_code):36    pyperclip.copy(malicious_code)37 38 39def activate_copy_button(o):40    return gr.Button.update(interactive=True)41    42# Create a Gradio interface43with gr.Blocks(title = "LLM Code injection", ) as demo:44    gr.Markdown("# LLM Code Injection Demo")45    gr.Markdown("This is a demonstration of how a malicious actor could use a prompt to inject code into a code generation model. The model will generate a form based on the prompt, but the prompt can be crafted to include malicious code that will be executed when the form is rendered and certain actions are taken. (For example, the submit button is pressed. The target of the attach would be casual users who use language models to generate code and run the code without verifying it)")46    # Add a noninteractive image47    #gr.Image(image_path, interactive=False, label="Demonstration Image")48 49    with gr.Row():50        safe_prompt_button = gr.Button("Enter safe prompt", min_width="100px")51        vulnerable_prompt_button = gr.Button("Enter vulnerable prompt", min_width="100px")52    53    # Add a text input and output interface54    with gr.Row():55        api_key_input = gr.Textbox(label="OpenAI API Key", type="password")56        text_input = gr.Textbox(label = "Prompt", placeholder = "Enter a prompt to generate a Svelte form")57        58    with gr.Row():    59        safe_code_output = gr.Textbox(label = "Code for form", info = "This is the code for the form that you can copy and render yourself by pasting it into a Svelte project.")60        html_output = gr.HTML(label="Svelte Form")61        malicious_code_output = gr.Textbox(label = "Malicious code", visible=False)62    63    with gr.Row():64        # Button to submit text65        submit_button = gr.Button("Submit")66        copy_button = gr.Button("Copy generated code")67    68    safe_prompt_button.click(generate_safe_prompt, None, outputs=[text_input])69    vulnerable_prompt_button.click(generate_vulnerable_prompt, None, outputs=[text_input])70    # Set the interaction between input and output71    submit_button.click(generate_crew_output, inputs=[text_input, api_key_input], outputs=[safe_code_output, html_output, malicious_code_output])72    copy_button.click(copy_malicious_code, inputs=[malicious_code_output]).then(raise_error)73 74demo.launch()75