Patched code sample

import io
import sys
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from RestrictedPython import compile_restricted, safe_builtins

# This code represents a hypothetical fix for a vulnerability like CVE-2025-51482.
# The original vulnerability would have likely used `exec()` or `eval()` on unsanitized
# user input, allowing arbitrary code execution.
#
# The fix replaces the direct, unsafe execution with a robust sandboxing mechanism
# using the 'RestrictedPython' library. This library compiles the user-provided code
# in a restricted mode, preventing access to dangerous modules, built-ins, and
# system-level functions.

app = FastAPI()

class ToolPayload(BaseModel):
    source_code: str
    tool_args: dict = {}

# Define a set of safe global variables that can be passed to the executed code.
# This prevents the code from accessing unsafe modules like 'os' or 'subprocess'.
# We can also add custom, safe functions here if needed.
restricted_globals = {
    "__builtins__": safe_builtins,
    # Example of a safe utility function that could be exposed to the tool
    "safe_print": print 
}

@app.post("/v1/tools/run")
async def run_tool_from_source(payload: ToolPayload):
    """
    Executes a tool from a source code string within a secure sandbox.
    """
    source_code = payload.source_code
    
    if not source_code.strip():
        raise HTTPException(status_code=400, detail="source_code cannot be empty.")

    # A dictionary to hold the local scope of the executed code,
    # including arguments and results.
    local_scope = {"args": payload.tool_args}
    
    # Use a string IO to capture any output from the sandboxed code
    output_buffer = io.StringIO()
    original_stdout = sys.stdout
    sys.stdout = output_buffer

    try:
        # 1. Compile the code in a restricted environment.
        #    This is the core of the security fix. `compile_restricted` will
        #    raise a SyntaxError if the code uses disallowed constructs
        #    (e.g., `import os`, opening files, accessing private attributes).
        byte_code = compile_restricted(
            source_code,
            filename='<user_tool>',
            mode='exec'
        )

        # 2. Execute the compiled bytecode.
        #    The execution is performed with the restricted globals and a
        #    specific local scope, preventing access to the wider application's state.
        exec(byte_code, restricted_globals, local_scope)

    except Exception as e:
        # If compilation or execution fails, restore stdout and raise an error.
        sys.stdout = original_stdout
        raise HTTPException(
            status_code=400,
            detail=f"Error executing tool: {e}"
        )
    finally:
        # Always restore the original stdout
        sys.stdout = original_stdout

    captured_output = output_buffer.getvalue()
    
    # The result of the operation is expected to be in the 'result'
    # variable within the executed code's local scope.
    result = local_scope.get("result", None)

    return {
        "status": "success",
        "result": result,
        "output": captured_output
    }

# Example of how to run this FastAPI application with uvicorn:
# uvicorn your_file_name:app --reload
#
# To test the fix:
#
# 1. Benign Payload (Should succeed):
# curl -X POST http://127.0.0.1:8000/v1/tools/run -H "Content-Type: application/json" -d \
# '{"source_code": "x = args.get(\"a\", 0) + args.get(\"b\", 0)\nresult = {\"sum\": x}", "tool_args": {"a": 5, "b": 10}}'
#
# 2. Malicious Payload (Should be blocked by RestrictedPython and fail):
# curl -X POST http://127.0.0.1:8000/v1/tools/run -H "Content-Type: application/json" -d \
# '{"source_code": "import os\nos.system(\"echo pwned > /tmp/pwned\")\nresult = \"failed\""}'
#
# The above request will result in an HTTP 400 error with a detail message like:
# "SyntaxError: ('Line 1: import os is not allowed.', ('<user_tool>', 1, 0, 'import os', 1, 10))"