Patched code sample

import uuid
from typing import Dict, Optional

from fastapi import APIRouter, Body, HTTPException, Path

# This is a mock database to simulate stored, trusted flow data.
# In a real application, this would be a persistent database (e.g., PostgreSQL, MySQL).
MOCK_DATABASE = {
    "public_flow_abc123": {
        "name": "Official Public Flow",
        "data": {
            "nodes": [
                {
                    "id": "node-1",
                    # This is the trusted code stored by the application owner.
                    "code": "print('Executing the safe, database-stored flow.')"
                }
            ]
        }
    }
}

router = APIRouter(prefix="/api/v1", tags=["Build"])


def get_flow_from_db(flow_id: str) -> Dict:
    """Simulates fetching a flow securely from the database."""
    flow = MOCK_DATABASE.get(flow_id)
    if not flow:
        raise HTTPException(status_code=404, detail="Flow not found in database.")
    return flow


def execute_flow_code(flow_data: Dict):
    """
    Simulates the part of the build process that executes code from a flow node.
    The `exec` call is the security-sensitive operation (sink).
    """
    try:
        code_to_run = flow_data["data"]["nodes"][0]["code"]
        print(f"Found trusted code to execute: '{code_to_run}'")
        # In a real fix, this 'exec' might be further sandboxed or replaced,
        # but the primary fix is ensuring the code source is trusted.
        exec(code_to_run, {}, {})
        print("Execution complete.")
    except (KeyError, IndexError) as e:
        print(f"Could not find code in the flow data structure: {e}")
    except Exception as e:
        print(f"An error occurred during code execution: {e}")


@router.post("/build_public_tmp/{flow_id}/flow", status_code=201)
def build_public_tmp_flow_fixed(
    flow_id: str = Path(..., description="The unique identifier for the public flow."),
    # The `data` parameter from the request body is still defined to handle incoming
    # requests gracefully, but its value is intentionally ignored.
    data: Optional[Dict] = Body(None, description="[DEPRECATED AND IGNORED] Flow data from the request body."),
):
    """
    Builds a public flow, demonstrating the fix for CVE-2026-33017.

    The fix ensures that any 'data' payload in the request body is
    **completely ignored**. The function exclusively uses the `flow_id` from the URL
    to fetch the flow configuration from the trusted database. This prevents an
    attacker from supplying arbitrary code in the request body for execution.
    """
    if data:
        # Log that extraneous data was received but is being ignored. This is a
        # critical part of the fix's logic.
        print(
            "INFO: Request body contained 'data', which is being ignored for security reasons "
            "on this public endpoint."
        )

    # --- CORE OF THE FIX ---
    # 1. The flow data is ALWAYS fetched from the database using the flow_id.
    #    The 'data' variable from the request body is never used.
    print(f"Fetching flow configuration for flow_id: {flow_id} from the database.")
    trusted_flow_data = get_flow_from_db(flow_id)

    # 2. The build and execution process now operates ONLY on the trusted data.
    print("Proceeding to build flow using only trusted data from the database.")
    execute_flow_code(trusted_flow_data)

    return {
        "message": "Flow build process initiated using trusted, stored data.",
        "flow_id": flow_id,
    }


# To run this example:
# 1. Install FastAPI and Uvicorn:
#    pip install "fastapi[all]"
# 2. Save the code as a Python file (e.g., `fixed_app.py`).
# 3. Run the server:
#    uvicorn fixed_app:app --reload
#
# To test the fix, send a POST request with malicious data:
#
# curl -X POST http://127.0.0.1:8000/api/v1/build_public_tmp/public_flow_abc123/flow \
# -H "Content-Type: application/json" \
# -d '{
#   "data": {
#     "nodes": [
#       {
#         "id": "malicious-node",
#         "code": "import os; os.system(\"echo VULNERABLE\")"
#       }
#     ]
#   }
# }'
#
# The server output will show that the malicious data was ignored and only the
# safe, database-stored code ("print(...)") was executed.
if __name__ == "__main__":
    import uvicorn
    from fastapi import FastAPI

    app = FastAPI()
    app.include_router(router)

    print("=" * 80)
    print("FastAPI server running. This demonstrates the fix for CVE-2026-33017.")
    print("Test by sending a POST request to /api/v1/build_public_tmp/{flow_id}/flow")
    print("Example using curl:")
    print(
        'curl -X POST "http://127.0.0.1:8000/api/v1/build_public_tmp/public_flow_abc123/flow" '
        '-H "Content-Type: application/json" '
        '-d \'{"data": {"nodes": [{"code": "import os; print(\'PWNED by attacker\')"}]}}\''
    )
    print("=" * 80)

    uvicorn.run(app, host="0.0.0.0", port=8000)