Nexus/server/api/rdp.py

"""Browser RDP — Guacamole WebSocket tunnel to guacd (no session audit)."""

from __future__ import annotations

import logging
from typing import Optional

from fastapi import APIRouter, Path, Query, WebSocket, WebSocketDisconnect

from server.config import settings
from server.domain.models import Admin
from server.infrastructure.database.session import AsyncSessionLocal
from server.infrastructure.database.admin_repo import AdminRepositoryImpl
from server.infrastructure.database.server_repo import ServerRepositoryImpl
from server.infrastructure.guacamole.tunnel import bridge_websocket_to_guacd
from server.utils.rdp_config import get_rdp_config

logger = logging.getLogger("nexus.rdp")

router = APIRouter()


async def _verify_rdp_access_token(token: str) -> Optional[Admin]:
    """Verify normal access JWT for RDP (reject webssh-only tokens)."""
    try:
        import jwt as pyjwt

        payload = pyjwt.decode(
            token,
            settings.SECRET_KEY,
            algorithms=["HS256"],
            options={"require": ["exp", "sub"]},
        )
    except Exception:
        return None

    if payload.get("purpose") == "webssh":
        return None

    admin_id = payload.get("sub")
    if not admin_id:
        return None
    try:
        admin_id = int(admin_id)
    except (TypeError, ValueError):
        return None

    async with AsyncSessionLocal() as session:
        admin_repo = AdminRepositoryImpl(session)
        admin = await admin_repo.get_by_id(admin_id)
        if not admin or not admin.is_active:
            return None
        token_tv = payload.get("tv") or 0
        if token_tv != (admin.token_version or 0):
            return None
        return admin


@router.websocket("/ws/rdp/{server_id}")
async def rdp_guacamole_ws(
    websocket: WebSocket,
    server_id: int = Path(...),
    token: Optional[str] = Query(None),
):
    """Transparent Guacamole tunnel to guacd after JWT + server RDP checks."""
    if not token:
        await websocket.close(code=4001, reason="Missing JWT token")
        return

    admin = await _verify_rdp_access_token(token)
    if not admin:
        await websocket.close(code=4401, reason="Invalid or expired JWT token")
        return

    if not settings.GUACD_HOST:
        await websocket.close(code=4503, reason="guacd not configured")
        return

    async with AsyncSessionLocal() as session:
        server_repo = ServerRepositoryImpl(session)
        server = await server_repo.get_by_id(server_id)
        if not server:
            await websocket.close(code=4404, reason="Server not found")
            return
        cfg = get_rdp_config(server)
        if not cfg.enabled:
            await websocket.close(code=4400, reason="RDP not enabled")
            return

    await websocket.accept(subprotocol="guacamole")
    try:
        await bridge_websocket_to_guacd(
            websocket,
            settings.GUACD_HOST,
            settings.GUACD_PORT,
        )
    except WebSocketDisconnect:
        pass
    except OSError as exc:
        logger.warning("RDP guacd bridge failed server_id=%s: %s", server_id, exc)
        try:
            await websocket.close(code=1011, reason="guacd unreachable")
        except Exception:
            pass
    except Exception:
        logger.exception("RDP tunnel error server_id=%s", server_id)
        try:
            await websocket.close(code=1011, reason="RDP tunnel error")
        except Exception:
            pass