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 urllib.parse import parse_qs, unquote

from fastapi import APIRouter, Path, WebSocket, WebSocketDisconnect

from server.config import settings
from server.domain.models import Admin
from server.infrastructure.database.crypto import decrypt_value
from server.infrastructure.database.session import AsyncSessionLocal
from server.infrastructure.database.admin_repo import AdminRepositoryImpl
from server.infrastructure.database.rdp_host_repo import RdpHostRepositoryImpl
from server.infrastructure.guacamole.ws_tunnel import run_guacamole_rdp_tunnel

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


def _rdp_password_plain(enc: str) -> str:
    plain = str(enc or "").strip()
    if not plain:
        return ""
    try:
        return decrypt_value(plain)
    except Exception:
        return ""


def _parse_display_size(query_string: bytes) -> tuple[int, int]:
    qs = parse_qs(query_string.decode("utf-8", errors="replace"))
    try:
        w = int(qs.get("width", ["1024"])[0])
        h = int(qs.get("height", ["768"])[0])
        if w > 0 and h > 0:
            return w, h
    except (TypeError, ValueError):
        pass
    return 1024, 768


@router.websocket("/ws/rdp/hosts/{host_id}/tunnel/{access_token}")
async def rdp_guacamole_ws(
    websocket: WebSocket,
    host_id: int = Path(...),
    access_token: str = Path(..., description="URL-encoded JWT (not query — Guacamole appends ?connect)"),
):
    """Guacamole WS tunnel: JWT auth, server-side guacd handshake, then relay."""
    token = unquote(access_token).strip()
    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:
        repo = RdpHostRepositoryImpl(session)
        row = await repo.get_by_id(host_id)
        if not row:
            await websocket.close(code=4404, reason="RDP host not found")
            return

        password = _rdp_password_plain(row.password)
        if not password:
            await websocket.close(code=4400, reason="RDP password not configured")
            return

        connect_params = {
            k: v
            for k, v in (
                ("hostname", row.host.strip()),
                ("port", str(row.port)),
                ("username", row.username.strip()),
                ("password", password),
                ("security", "any"),
                ("ignore-cert", "true"),
                ("disable-wallpaper", "true"),
                ("disable-theming", "true"),
                ("enable-font-smoothing", "true"),
                ("enable-full-window-drag", "false"),
                ("enable-desktop-composition", "false"),
                ("enable-menu-animations", "false"),
            )
            if v
        }

    width, height = _parse_display_size(websocket.scope.get("query_string", b""))

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