Vulnerable code sample

import asyncio
import logging
import subprocess
import sys

from starlette.websockets import WebSocket

from marimo._server.model import SessionMode
from marimo._server.router import APIRouter
from marimo._server.sessions import SessionManager

LOGGER = logging.getLogger(__name__)

# Router for terminal-related endpoints
router = APIRouter()


# This is a hypothetical representation of the server's session manager
# It is used to check the server's running mode.
session_manager = SessionManager(...)


@router.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    # This is a placeholder for how other endpoints correctly
    # performed authentication. The real implementation would use
    # a dependency injection system like FastAPI's `Depends`.
    # auth_token: str = Depends(validate_auth),
):
    # This endpoint is assumed to be secure and is for contrast.
    # The actual implementation would validate auth before this point.
    await websocket.accept()
    await websocket.send_text("Authenticated and connected.")
    await websocket.close()


@router.websocket("/terminal/ws")
async def terminal_ws(
    websocket: WebSocket,
) -> None:
    """
    This endpoint provides a websocket connection to a terminal.

    VULNERABILITY:
    This handler does NOT perform any authentication checks.
    It only checks the server mode and platform, allowing any unauthenticated
    user to connect and execute arbitrary commands.
    """
    # Check if the terminal is allowed in the current mode
    if session_manager.mode != SessionMode.EDIT:
        LOGGER.debug("Terminal disabled in mode %s", session_manager.mode)
        await websocket.close(code=1008)
        return

    # Check for platform support
    if sys.platform == "win32":
        LOGGER.debug("Terminals are not supported on Windows.")
        await websocket.close(code=1003)
        return

    # No authentication is performed before accepting the connection.
    # An attacker can connect directly to this endpoint.
    await websocket.accept()

    # The following code represents how an attacker could gain a PTY shell.
    # This is a simplified representation of creating and interacting with a
    # a shell process.
    try:
        # Create a shell process
        process = await asyncio.create_subprocess_shell(
            "/bin/bash" if sys.platform != "win32" else "cmd.exe",
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )

        async def read_from_ws():
            # Forward commands from the websocket to the shell's stdin
            async for message in websocket.iter_text():
                if process.stdin:
                    process.stdin.write(message.encode())
                    await process.stdin.drain()

        async def write_to_ws():
            # Forward output from the shell's stdout/stderr to the websocket
            while True:
                output = await process.stdout.read(1)
                if not output:
                    break
                await websocket.send_bytes(output)

        # Run both tasks concurrently
        read_task = asyncio.create_task(read_from_ws())
        write_task = asyncio.create_task(write_to_ws())
        await asyncio.gather(read_task, write_task)

    except Exception as e:
        LOGGER.error(f"Terminal session error: {e}")
    finally:
        if "process" in locals() and process.returncode is None:
            process.kill()
        await websocket.close()