homelab_automation/app/routes/websocket.py

"""
Routes WebSocket pour les mises à jour en temps réel.
"""

import asyncio
import logging
from typing import Optional

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, HTTPException, status

from app.core.config import settings
from app.models.database import async_session_maker
from app.services import ws_manager
from app.crud.terminal_session import TerminalSessionRepository
from app.services.terminal_service import terminal_service

logger = logging.getLogger(__name__)

router = APIRouter()


@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """Endpoint WebSocket pour les mises à jour en temps réel."""
    await ws_manager.connect(websocket)
    try:
        while True:
            # Garder la connexion ouverte
            data = await websocket.receive_text()
            # Traiter les messages entrants si nécessaire
    except WebSocketDisconnect:
        ws_manager.disconnect(websocket)


@router.websocket("/terminal/ws/{session_id}")
async def terminal_websocket_proxy(
    websocket: WebSocket,
    session_id: str,
):
    """
    WebSocket proxy for ttyd terminal sessions.
    
    This endpoint proxies WebSocket connections to the local ttyd instance,
    allowing the browser to connect through the exposed port 8008 instead of
    needing direct access to the ttyd port (7680+).
    
    The session_id is used to look up the ttyd port from the database.
    Authentication is handled via session token in the query string.
    """
    # Get token from query params
    token = websocket.query_params.get("token")
    if not token:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="Missing token")
        return
    
    async with async_session_maker() as db_session:
        session_repo = TerminalSessionRepository(db_session)
        session = await session_repo.get(session_id)
        
        if not session:
            await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="Session not found")
            return
        
        # Verify token
        if not terminal_service.verify_token(token, session.token_hash):
            await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="Invalid token")
            return
    
    # Accept the WebSocket connection
    await websocket.accept()
    
    # Connect to local ttyd via TCP socket
    ttyd_port = session.ttyd_port
    ttyd_host = "127.0.0.1"  # ttyd is always local in the container
    
    try:
        # Create a TCP socket to ttyd
        reader, writer = await asyncio.open_connection(ttyd_host, ttyd_port)
        logger.info(f"Proxying WebSocket for session {session_id[:8]}... to ttyd on port {ttyd_port}")
        
        # Bidirectional proxy: forward messages between WebSocket and ttyd
        async def forward_ws_to_ttyd():
            """Forward WebSocket messages to ttyd."""
            try:
                while True:
                    data = await websocket.receive_bytes()
                    writer.write(data)
                    await writer.drain()
            except WebSocketDisconnect:
                logger.debug(f"WebSocket disconnected for session {session_id[:8]}...")
            except Exception as e:
                logger.error(f"Error forwarding WS to ttyd: {e}")
            finally:
                try:
                    writer.close()
                    await writer.wait_closed()
                except Exception:
                    pass
        
        async def forward_ttyd_to_ws():
            """Forward ttyd messages to WebSocket."""
            try:
                while True:
                    data = await reader.readexactly(1024)
                    if not data:
                        break
                    await websocket.send_bytes(data)
            except asyncio.IncompleteReadError:
                # Normal EOF
                pass
            except Exception as e:
                logger.error(f"Error forwarding ttyd to WS: {e}")
            finally:
                try:
                    await websocket.close()
                except Exception:
                    pass
        
        # Run both directions concurrently
        await asyncio.gather(
            forward_ws_to_ttyd(),
            forward_ttyd_to_ws(),
            return_exceptions=True
        )
        
    except ConnectionRefusedError:
        logger.error(f"Failed to connect to ttyd on {ttyd_host}:{ttyd_port} for session {session_id[:8]}...")
        await websocket.close(code=1011, reason="ttyd connection failed")
    except Exception as e:
        logger.exception(f"Error in terminal WebSocket proxy: {e}")
        try:
            await websocket.close(code=1011, reason="Proxy error")
        except Exception:
            pass