Patched code sample

from fastapi import FastAPI, Request, HTTPException, Security, Depends
from fastapi.security.api_key import APIKeyHeader
import subprocess
import tempfile
import os

app = FastAPI()

# In a real application, this secret should be managed via environment variables or a secrets manager.
VALID_API_KEY = os.environ.get("LANGFLOW_API_KEY", "SECURE_DEFAULT_KEY_CHANGE_ME")

api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)

# FIX 1: Implement a mandatory authentication dependency.
# This function enforces that a valid API key must be present in the request headers,
# fixing the authentication bypass vulnerability.
async def require_authentication(api_key: str = Security(api_key_header)):
    if api_key is None or api_key != VALID_API_KEY:
        raise HTTPException(
            status_code=401, detail="Unauthorized: Invalid or missing API Key"
        )
    return True

# FIX 2: Replace direct, un-sandboxed code execution (like exec()) with an isolated,
# controlled execution environment.
def run_code_in_isolated_sandbox(code: str) -> dict:
    """
    Executes Python code inside a secure, isolated environment.

    This is a conceptual fix. A robust implementation would use dedicated sandboxing
    technologies like Docker, gVisor, or Firecracker to create an ephemeral,
    unprivileged container for each execution. This simplified example uses
    `subprocess` with security flags and resource limits as a demonstration.
    """
    # Use a temporary file to avoid command-line injection vulnerabilities.
    with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as tmp_file:
        tmp_file.write(code)
        script_path = tmp_file.name

    try:
        # Execute the Python script in a new, isolated process with strict limits.
        # -I flag isolates Python from the user's environment.
        # prlimit sets resource limits (e.g., max 1 process, 5-second CPU time).
        # This prevents fork bombs and long-running processes.
        # Network access should be disabled at a lower level (e.g., namespaces, firewall).
        result = subprocess.run(
            [
                "prlimit", "--nproc=1", "--cpu=5",
                "python3", "-I", script_path
            ],
            capture_output=True,
            text=True,
            timeout=10,  # A hard timeout for the subprocess call itself.
            check=False,
        )
        return {
            "stdout": result.stdout,
            "stderr": result.stderr,
            "returncode": result.returncode,
        }
    except subprocess.TimeoutExpired:
        return {
            "stdout": "",
            "stderr": "Execution timed out.",
            "returncode": -1,
        }
    finally:
        # Ensure the temporary script is always removed.
        os.remove(script_path)


# The fixed API endpoint now applies the authentication dependency and uses the sandbox.
@app.post("/api/v1/run_code", dependencies=[Depends(require_authentication)])
async def run_code_securely_endpoint(request: Request):
    """
    An endpoint that securely executes user-provided Python code.
    """
    try:
        data = await request.json()
        code_to_run = data.get("code")
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid JSON body.")

    if not isinstance(code_to_run, str) or not code_to_run.strip():
        raise HTTPException(status_code=400, detail="`code` field must be a non-empty string.")

    # Execute the code using the sandboxed function instead of a direct 'exec()'.
    execution_result = run_code_in_isolated_sandbox(code_to_run)

    return {
        "message": "Execution finished.",
        "result": execution_result
    }