FastAPI sever to execute unix commands and stream back the response via websockets

main.py

# main.py

import asyncio
from fastapi import FastAPI, WebSocket, WebSocketDisconnect

app = FastAPI()


async def stream_shell_output(command: str, websocket: WebSocket) -> None:
    """
    Execute a shell command and stream the output to the WebSocket client.
    """
    process = await asyncio.create_subprocess_shell(
        command,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT,
    )

    try:
        assert process.stdout is not None
        while True:
            line = await process.stdout.readline()
            if not line:
                break
            await websocket.send_text(line.decode().rstrip())

        await process.wait()
        await websocket.send_text(f"[EXIT CODE] {process.returncode}")
    except asyncio.CancelledError:
        process.kill()
        await websocket.send_text("[CANCELLED]")
        raise
    finally:
        await websocket.close()


@app.websocket("/ws/command")
async def websocket_command_handler(websocket: WebSocket):
    await websocket.accept()
    try:
        command: str = await websocket.receive_text()
        await stream_shell_output(command, websocket)
    except WebSocketDisconnect:
        print("WebSocket client disconnected")
    except Exception as e:
        await websocket.send_text(f"[ERROR] {str(e)}")
        await websocket.close()


if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", reload=True)