kvm-manager/frontend/serve.py

"""KVM Frontend - Static files + API proxy + WebSocket VNC proxy"""
import os
import re

from starlette.applications import Starlette
from starlette.routing import Route, WebSocketRoute, Mount
from starlette.staticfiles import StaticFiles
from starlette.responses import FileResponse, RedirectResponse
from starlette.websockets import WebSocket
import uvicorn

HOST = "0.0.0.0"
PORT = 8006
DIST_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "dist")
API_BASE = "http://127.0.0.1:8004"
WS_BASE = "ws://127.0.0.1:8004"


# ── SPA fallback ──────────────────────────────────────────────
async def serve_spa(request):
    path = request.url.path
    file_path = os.path.join(DIST_DIR, path.lstrip("/"))
    if os.path.isfile(file_path):
        return FileResponse(file_path)
    return FileResponse(os.path.join(DIST_DIR, "index.html"))


# ── API proxy ─────────────────────────────────────────────────
async def proxy_api(request):
    from starlette.requests import Request
    body = await request.body()
    headers = {}
    for key in ["content-type", "authorization"]:
        val = request.headers.get(key)
        if val:
            headers[key.replace("_", "-")] = val

    url = f"{API_BASE}{request.url.path}"
    if request.url.query:
        url += f"?{request.url.query}"

    async with __import__("httpx").AsyncClient(timeout=30) as client:
        try:
            resp = await client.request(
                method=request.method,
                url=url,
                content=body,
                headers=headers,
            )
            return Response(
                resp.content,
                status_code=resp.status_code,
                headers={"access-control-allow-origin": "*"},
            )
        except Exception as e:
            from starlette.responses import JSONResponse
            return JSONResponse({"error": str(e)}, status_code=502)
    from starlette.responses import Response
    return Response("", status_code=500)


# ── WebSocket VNC proxy ────────────────────────────────────────
async def vnc_websocket(websocket: WebSocket):
    await websocket.accept()

    # Build backend WebSocket URL
    path = websocket.url.path  # e.g. /ws/vnc/vm-name?host_id=local
    backend_url = f"{WS_BASE}{path}"

    try:
        async with __import__("websockets").connect(backend_url) as backend:
            async def forward_to_backend():
                async for msg in websocket.iter_text():
                    await backend.send(msg)
                async for msg in websocket.iter_bytes():
                    await backend.send(msg)

            async def forward_to_frontend():
                async for msg in backend:
                    if isinstance(msg, bytes):
                        await websocket.send_bytes(msg)
                    else:
                        await websocket.send_text(msg)

            await __import__("asyncio").gather(
                forward_to_backend(),
                forward_to_frontend(),
            )
    except Exception as e:
        print(f"WebSocket proxy error: {e}")
    finally:
        try:
            await websocket.close()
        except Exception:
            pass


from starlette.responses import Response

routes = [
    WebSocketRoute("/ws/{path:path}", vnc_websocket, name="vnc"),
    Route("/api/{path:path}", proxy_api),
    Route("/api", proxy_api),
    Mount("/", StaticFiles(directory=DIST_DIR, html=True), name="static"),
]

app = Starlette(routes=routes)

if __name__ == "__main__":
    uvicorn.run(
        app,
        host=HOST,
        port=PORT,
        log_level="info",
    )