guardia-itsm/routers/ws.py

"""
GUARDiA ITSM — WebSocket 실시간 이벤트 푸시 (Enhancement A-1)

기능:
  1. WebSocket 연결 관리 (ConnectionManager)
  2. JWT 토큰 기반 인증 (query parameter ?token=...)
  3. 역할별 구독 채널 필터링
  4. SSE broadcast() 와 동일한 이벤트 수신
  5. 하트비트 (30초 간격 ping)
  6. 재연결 지원 (last_event_id 기반)

엔드포인트:
  WS /ws/events?token={jwt}          — 전체 이벤트 스트림
  WS /ws/events/{channel}?token={jwt} — 채널별 (sr, deploy, sla, oncall, batch)
  GET /api/ws/status                 — WebSocket 연결 상태 (ADMIN)

프론트엔드 연결 예시:
  const ws = new WebSocket("ws://localhost:8000/ws/events?token=" + accessToken);
  ws.onmessage = (evt) => {
      const msg = JSON.parse(evt.data);
      if (msg.type === "sr_created") { ... }
  };
"""
from __future__ import annotations

import asyncio
import json
import logging
from datetime import datetime
from typing import Dict, Optional, Set

from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect, Query, HTTPException
from sqlalchemy.ext.asyncio import AsyncSession

from core.auth import get_current_user
from database import get_db
from models import User, UserRole

logger = logging.getLogger(__name__)
router = APIRouter(tags=["websocket"])

# ── 채널 정의 ─────────────────────────────────────────────────────────────────
_CHANNELS = {"sr", "deploy", "sla", "oncall", "batch", "all"}
_CHANNEL_EVENT_MAP: Dict[str, Set[str]] = {
    "sr":     {"sr_created", "sr_updated", "sr_status_changed"},
    "deploy": {"deploy_started", "deploy_completed", "deploy_failed", "vibe_updated"},
    "sla":    {"sla_violation", "sla_escalated", "sr_updated"},
    "oncall": {"oncall_assigned", "oncall_escalated"},
    "batch":  {"batch_started", "batch_completed", "batch_failed"},
    "all":    None,    # None = 모든 이벤트
}

_HEARTBEAT_INTERVAL = 30   # 초


# ═══════════════════════════════════════════════════════════════════════════════
# ── 연결 관리자 ────────────────────────────────────────────────────────────────
# ═══════════════════════════════════════════════════════════════════════════════

class ConnectionManager:
    """WebSocket 연결 목록 관리."""

    def __init__(self):
        # { websocket: { "username": str, "role": str, "channel": str, "connected_at": datetime } }
        self._connections: Dict[WebSocket, dict] = {}

    async def connect(
        self,
        ws: WebSocket,
        username: str,
        role: str,
        channel: str = "all",
    ) -> None:
        await ws.accept()
        self._connections[ws] = {
            "username":     username,
            "role":         role,
            "channel":      channel,
            "connected_at": datetime.now().isoformat(),
        }
        logger.info("WS 연결: user=%s channel=%s total=%d", username, channel, len(self._connections))

    def disconnect(self, ws: WebSocket) -> None:
        info = self._connections.pop(ws, {})
        logger.info(
            "WS 연결 해제: user=%s total=%d",
            info.get("username", "unknown"),
            len(self._connections),
        )

    async def broadcast(self, event_type: str, data: dict) -> None:
        """특정 이벤트 타입을 구독 중인 모든 WebSocket에 전송."""
        if not self._connections:
            return
        msg = json.dumps({"type": event_type, **data}, ensure_ascii=False)
        dead: list[WebSocket] = []

        for ws, info in list(self._connections.items()):
            channel = info.get("channel", "all")
            # 채널 필터링
            allowed = _CHANNEL_EVENT_MAP.get(channel)
            if allowed is not None and event_type not in allowed:
                continue
            # CUSTOMER는 자신의 기관 이벤트만 수신 (향후 확장)
            try:
                await ws.send_text(msg)
            except Exception:
                dead.append(ws)

        for ws in dead:
            self.disconnect(ws)

    def connection_count(self) -> int:
        return len(self._connections)

    def connections_info(self) -> list:
        return [
            {
                "username":     info["username"],
                "role":         info["role"],
                "channel":      info["channel"],
                "connected_at": info["connected_at"],
            }
            for info in self._connections.values()
        ]

    async def send_personal(self, ws: WebSocket, message: dict) -> bool:
        """특정 WebSocket에만 메시지 전송."""
        try:
            await ws.send_text(json.dumps(message, ensure_ascii=False))
            return True
        except Exception:
            self.disconnect(ws)
            return False


# 전역 연결 관리자 (앱 수명 동안 유지)
manager = ConnectionManager()


# ── core/events.py 와 통합 ────────────────────────────────────────────────────

def _integrate_with_sse_bus() -> None:
    """
    core/events.broadcast() 호출 시 WebSocket manager.broadcast() 도 함께 실행되도록
    이벤트 버스를 패치합니다.

    main.py lifespan 에서 한 번 호출됩니다.
    """
    import core.events as _events

    _original_broadcast = _events.broadcast

    async def _patched_broadcast(event_type: str, data=None):
        # 기존 SSE 브로드캐스트 실행
        await _original_broadcast(event_type, data)
        # WebSocket 브로드캐스트도 실행
        await manager.broadcast(event_type, data or {})

    _events.broadcast = _patched_broadcast
    logger.info("WebSocket 통합: core/events.broadcast() 패치 완료")


# ═══════════════════════════════════════════════════════════════════════════════
# ── WebSocket 인증 헬퍼 ────────────────────────────────────────────────────────
# ═══════════════════════════════════════════════════════════════════════════════

async def _authenticate_ws(token: str, db: AsyncSession) -> Optional[User]:
    """WebSocket 연결 시 JWT 토큰으로 사용자 인증."""
    if not token:
        return None
    try:
        from core.auth import SECRET_KEY, ALGORITHM
        from jose import jwt, JWTError
        from sqlalchemy import select

        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        # mfa_pending 토큰은 거부
        if payload.get("mfa_pending"):
            return None
        username = payload.get("sub")
        if not username:
            return None

        result = await db.execute(select(User).where(User.username == username))
        user = result.scalars().first()
        return user if (user and user.is_active) else None
    except Exception:
        return None


# ═══════════════════════════════════════════════════════════════════════════════
# ── WebSocket 엔드포인트 ──────────────────────────────────────────────────────
# ═══════════════════════════════════════════════════════════════════════════════

@router.websocket("/ws/events")
async def ws_all_events(
    websocket: WebSocket,
    token: str = Query(..., description="JWT access_token"),
    db: AsyncSession = Depends(get_db),
):
    """전체 이벤트 스트림 WebSocket."""
    await _handle_ws(websocket, token, "all", db)


@router.websocket("/ws/events/{channel}")
async def ws_channel_events(
    websocket: WebSocket,
    channel: str,
    token: str = Query(..., description="JWT access_token"),
    db: AsyncSession = Depends(get_db),
):
    """채널별 이벤트 WebSocket. channel: sr|deploy|sla|oncall|batch|all"""
    if channel not in _CHANNELS:
        await websocket.close(code=4001, reason=f"Unknown channel: {channel}")
        return
    await _handle_ws(websocket, token, channel, db)


async def _handle_ws(
    websocket: WebSocket,
    token: str,
    channel: str,
    db: AsyncSession,
) -> None:
    """WebSocket 연결 처리 공통 로직."""
    # 인증
    user = await _authenticate_ws(token, db)
    if not user:
        await websocket.close(code=4001, reason="인증 실패: 유효한 토큰이 필요합니다.")
        return

    await manager.connect(websocket, user.username, user.role, channel)

    # 연결 성공 메시지
    await manager.send_personal(websocket, {
        "type":         "connected",
        "username":     user.username,
        "channel":      channel,
        "server_time":  datetime.now().isoformat(),
        "message":      f"GUARDiA ITSM WebSocket 연결됨 (채널: {channel})",
    })

    try:
        # 하트비트 태스크
        heartbeat_task = asyncio.create_task(_heartbeat_loop(websocket))

        # 클라이언트 메시지 수신 루프 (ping/pong, 구독 변경 등)
        while True:
            try:
                raw = await asyncio.wait_for(websocket.receive_text(), timeout=_HEARTBEAT_INTERVAL + 5)
                msg = json.loads(raw)
                await _handle_client_message(websocket, msg, user)
            except asyncio.TimeoutError:
                # 클라이언트가 아무것도 안 보내면 그냥 계속
                pass

    except WebSocketDisconnect:
        pass
    except Exception as exc:
        logger.debug("WS 오류: user=%s err=%s", user.username, exc)
    finally:
        heartbeat_task.cancel()
        manager.disconnect(websocket)


async def _heartbeat_loop(ws: WebSocket) -> None:
    """30초마다 서버 heartbeat 전송."""
    while True:
        await asyncio.sleep(_HEARTBEAT_INTERVAL)
        try:
            await ws.send_text(json.dumps({
                "type":        "heartbeat",
                "server_time": datetime.now().isoformat(),
                "connections": manager.connection_count(),
            }, ensure_ascii=False))
        except Exception:
            break


async def _handle_client_message(ws: WebSocket, msg: dict, user: User) -> None:
    """클라이언트에서 보낸 메시지 처리."""
    msg_type = msg.get("type", "")

    if msg_type == "ping":
        await manager.send_personal(ws, {
            "type": "pong",
            "server_time": datetime.now().isoformat(),
        })
    elif msg_type == "subscribe":
        new_channel = msg.get("channel", "all")
        if new_channel in _CHANNELS:
            info = manager._connections.get(ws)
            if info:
                info["channel"] = new_channel
            await manager.send_personal(ws, {
                "type": "subscribed",
                "channel": new_channel,
            })
    elif msg_type == "status":
        await manager.send_personal(ws, {
            "type":        "status_reply",
            "connections": manager.connection_count(),
            "username":    user.username,
            "role":        user.role,
        })


# ═══════════════════════════════════════════════════════════════════════════════
# ── REST API: WebSocket 연결 상태 조회 ───────────────────────────────────────
# ═══════════════════════════════════════════════════════════════════════════════

@router.get("/api/ws/status")
async def ws_status(
    current_user: User = Depends(get_current_user),
):
    """WebSocket 연결 상태 조회 (ADMIN 전용)."""
    if current_user.role != UserRole.ADMIN:
        raise HTTPException(403, "ADMIN 권한이 필요합니다.")

    return {
        "total_connections":  manager.connection_count(),
        "connections":        manager.connections_info(),
        "channels":           list(_CHANNELS),
    }


@router.post("/api/ws/broadcast")
async def ws_broadcast_manual(
    event_type: str = Query(..., description="이벤트 타입"),
    message: str    = Query("", description="메시지"),
    current_user: User = Depends(get_current_user),
):
    """수동 WebSocket 브로드캐스트 (ADMIN 전용 — 테스트/공지용)."""
    if current_user.role != UserRole.ADMIN:
        raise HTTPException(403, "ADMIN 권한이 필요합니다.")

    await manager.broadcast(event_type, {
        "message":    message,
        "sent_by":    current_user.username,
        "sent_at":    datetime.now().isoformat(),
    })
    return {
        "sent_to":    manager.connection_count(),
        "event_type": event_type,
    }