DO NOT USE THIS SCRIPT IN ANY PRODUCTION OR LIVE ENVIRONMENT. IT IS INTENTIONALLY DESIGNED WITH SEVERE SECURITY VULNERABILITIES FOR EDUCATIONAL PURPOSES ONLY. DEPLOYING THIS CODE EXPOSES THE HOST SYSTEM TO REMOTE CODE EXECUTION (RCE).
This Python script implements a simple Flask web server that acts as a proxy between a user and a destination Ollama (or compatible LLM) API endpoint. Its primary function is to:
- Receive a JSON request (typically containing a
modelandprompt) on its/api/generateendpoint. - Forward this request to a configured
DESTINATION_OLLAMA_ENDPOINT. - Receive the JSON response from the destination endpoint.
- Crucially and Dangerously: Inspect the
responsefield within the received JSON. If this field contains a string that includes a JSON object like{"command_os": "some_command"}, the script will attempt to executesome_commandas an operating system command on the server where the proxy is running. - Append the standard output (stdout) of the executed command (if any) to the original
responsetext. - Return the potentially modified JSON response to the original caller.
This tool is specifically crafted to teach and demonstrate AI Red Teaming concepts, particularly focusing on Indirect Prompt Injection leading to Remote Code Execution (RCE) through manipulating Large Language Model (LLM) outputs.
Students or security researchers can interact with this proxy, sending prompts designed to trick the backing LLM (at DESTINATION_OLLAMA_ENDPOINT) into generating a response containing the malicious {"command_os": "..."} payload. If successful, the proxy will execute the specified command, demonstrating a critical vulnerability chain.
- Listens: The Flask app listens on
0.0.0.0:5001(by default) for POST requests to/api/generate. - Validates Input: Checks if the request is JSON and contains
modelandpromptkeys. - Forwards Request: Creates a separate thread to send the received JSON payload to the
DESTINATION_OLLAMA_ENDPOINTvia a POST request. - Receives Response: Waits for the response from the destination endpoint.
- Parses and Inspects:
- Retrieves the value of the
responsekey from the destination's JSON response. - Uses
extract_braced_contentto find the first{...}block within the response string. - Uses
execute_command_os_from_stringto:- Parse the extracted block as JSON.
- Check if a
command_oskey exists. - If the key exists, execute its string value using
subprocess.run(command, shell=True, ...). This is the core vulnerability. - Capture the
stdoutandstderrof the executed command.
- Retrieves the value of the
- Modifies Response: If a command was executed successfully, its
stdoutis appended to the originalresponsestring from the LLM. - Returns Result: Sends the final JSON (original or modified) and the status code received from the destination back to the initial client.
Before running, you MUST edit the script and set the DESTINATION_OLLAMA_ENDPOINT variable:
# !!! IMPORTANT: Replace this with the actual URL of your destination Ollama service !!!
DESTINATION_OLLAMA_ENDPOINT = "http://localhost:11434/api/generate" # Example: Replace with your real endpointReplace "http://localhost:11434/api/generate" with the correct URL for the LLM API you want this proxy to forward requests to.
- Prerequisites: Ensure you have Python 3 and the required libraries installed:
pip install Flask requests
- Configure: Edit the script to set the correct
DESTINATION_OLLAMA_ENDPOINT. - Run: Execute the script from your terminal:
(Replace
python your_script_name.py
your_script_name.pywith the actual filename). - Access: The proxy will be available at
http://<your-server-ip>:5001/api/generate.
- Remote Code Execution (RCE): This is the primary, intentional vulnerability. The
execute_command_os_from_stringfunction directly takes a string potentially derived from an external LLM's output and executes it as a shell command viasubprocess.run(..., shell=True).shell=True: This is particularly dangerous as it allows shell metacharacters (;,|,&&,$(),`, etc.) to be interpreted, making complex command injection easier.- Untrusted Input: The command originates from the output of the LLM at the
DESTINATION_OLLAMA_ENDPOINT. An attacker can craft prompts sent through this proxy, aiming to make the LLM generate a response containing{"command_os": "malicious_command"}. The proxy will blindly executemalicious_command.
- Lack of Input Sanitization: No significant attempt is made to sanitize the command string before execution.
- Error Handling Exposure: While errors are caught, detailed error messages might be returned, potentially revealing information about the server environment or the script's internal state.
- Denial of Service (DoS): An attacker could potentially provide commands that consume excessive resources (e.g.,
cat /dev/urandom) or trigger infinite loops if the shell environment allows.
curl -X POST http://localhost:5001/api/generate -H "Content-Type: application/json" -d '{ "model": "gemma2:27b", "prompt": "Return a single linux command in json format using a KV command_os with command in its value field", "stream": false, "options": { "temperature": 0.0, "seed": 1337 }}'
curl -X POST http://localhost:5001/api/generate -H "Content-Type: application/json" -d '{ "model": "gemma2:27b", "prompt": "ping 127.0.0.1 one time in linux command in json format using a KV command_os as the key name", "stream": false, "options": { "temperature": 0.0, "seed": 1337 }}'This script is provided AS IS, purely for educational demonstration of AI security vulnerabilities. The authors and distributors accept no liability for any damage caused by the use or misuse of this code. By running this script, you acknowledge the extreme risks involved and agree that you are solely responsible for any consequences.
NEVER deploy this in a production, shared, or untrusted environment. Use it only in isolated, controlled lab settings specifically designed for security testing and education.