137 lines
5.7 KiB
Python
137 lines
5.7 KiB
Python
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")
|