services

app/api/ws.py

@@ -0,0 +1,136 @@
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Query
+from typing import List, Dict
+import asyncio
+import logging
+import json
+from sqlalchemy.orm import Session
+from app.database import SessionLocal
+from app.api.auth import get_current_user_ws
+from app.schemas.auth import CurrentUser
+from app.services.proxmox_service import proxmox_service
+from app.models.vm import VMAccess
+from app.schemas.vm import VMInfo
+
+router = APIRouter()
+logger = logging.getLogger(__name__)
+
+class ConnectionManager:
+    def __init__(self):
+        # Active connections: List of (WebSocket, User) tuples
+        self.active_connections: List[tuple[WebSocket, CurrentUser]] = []
+
+    async def connect(self, websocket: WebSocket, user: CurrentUser):
+        await websocket.accept()
+        self.active_connections.append((websocket, user))
+        logger.info(f"WebSocket connected: {user.username}")
+
+    def disconnect(self, websocket: WebSocket, user: CurrentUser):
+        if (websocket, user) in self.active_connections:
+            self.active_connections.remove((websocket, user))
+            logger.info(f"WebSocket disconnected: {user.username}")
+
+    async def broadcast(self, all_resources: List[Dict]):
+        """
+        모든 연결된 클라이언트에게 권한에 맞는 VM 상태를 전송
+        """
+        # DB 세션을 매번 새로 생성하는 것은 비효율적일 수 있으나, 
+        # Background Task에서 실행되므로 안전하게 처리
+        try:
+            db = SessionLocal()
+            
+            # 모든 활성 VM Access 정보 미리 로딩
+            all_accesses = db.query(VMAccess).filter(VMAccess.is_active == True).all()
+            
+            # User ID별 Access Map 생성
+            user_access_map = {}
+            for access in all_accesses:
+                if access.user_id not in user_access_map:
+                    user_access_map[access.user_id] = {}
+                user_access_map[access.user_id][access.vm_id] = access
+            
+            for connection, user in self.active_connections:
+                try:
+                    # 해당 유저의 권한 맵 가져오기
+                    access_map = user_access_map.get(user.id, {})
+                    
+                    user_vm_list = []
+                    for res in all_resources:
+                        vm_id = res.get("vmid")
+                        # VMID가 없거나 권한 정보가 없으면 패스
+                        if not vm_id or vm_id not in access_map:
+                            continue
+                        
+                        access = access_map[vm_id]
+                        
+                        # VMInfo 스키마에 맞춰 데이터 구성
+                        vm_info = {
+                            "vm_id": vm_id,
+                            "node": res.get("node"),
+                            "type": res.get("type", "qemu"),
+                            "name": res.get("name", "Unknown"),
+                            "status": res.get("status", "unknown"),
+                            "ip_address": access.static_ip,
+                            "cpus": res.get("maxcpu", 0),
+                            "memory": res.get("maxmem", 0) // (1024 * 1024),
+                            "memory_usage": res.get("mem", 0) // (1024 * 1024) if res.get("mem") else 0,
+                            "cpu_usage": res.get("cpu", 0),
+                            # 권한 정보
+                            "can_start": access.can_start,
+                            "can_stop": access.can_stop,
+                            "can_reboot": access.can_reboot,
+                            "can_connect": access.can_connect,
+                            # RDP 정보는 보안상 웹소켓 브로드캐스트에서는 제외하거나 필요시 포함 
+                            # (여기서는 제외하고 REST API 상세조회 사용 권장하지만, 목록 뷰를 위해 포함)
+                            "rdp_username": access.rdp_username,
+                            "rdp_port": access.rdp_port or 3389
+                        }
+                        user_vm_list.append(vm_info)
+                    
+                    # 전송
+                    await connection.send_json({
+                        "type": "update",
+                        "data": user_vm_list
+                    })
+                    
+                except Exception as e:
+                    logger.error(f"Error sending update to {user.username}: {e}")
+                    # 에러 발생 시 연결 끊기 고려?
+            
+            db.close()
+            
+        except Exception as e:
+            logger.error(f"Broadcast error: {e}")
+
+manager = ConnectionManager()
+
+@router.websocket("/status")
+async def websocket_endpoint(
+    websocket: WebSocket,
+    user: CurrentUser = Depends(get_current_user_ws)
+):
+    try:
+        await manager.connect(websocket, user)
+        
+        try:
+            while True:
+                # 클라이언트로부터 메시지를 받을 일이 딱히 없지만 연결 유지를 위해 대기
+                await websocket.receive_text()
+        except WebSocketDisconnect:
+            manager.disconnect(websocket, user)
+            
+    except Exception:
+        await websocket.close(code=4001)
+
+async def start_monitoring_task():
+    """백그라운드 모니터링 태스크"""
+    logger.info("Starting background monitoring task...")
+    try:
+        while True:
+            # 활성 연결이 있을 때만 조회 (리소스 절약)
+            if manager.active_connections:
+                resources = await proxmox_service.get_all_vms()
+                await manager.broadcast(resources)
+            
+            await asyncio.sleep(2)  # 2초마다 갱신
+    except asyncio.CancelledError:
+        logger.info("Monitoring task cancelled")