kvm-manager/backend/app/main.py

"""FastAPI 主应用"""
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from app.config import settings
from app.routers import vm, storage, network, snapshot, monitor, auth, host as host_router

app = FastAPI(
    title=settings.APP_NAME,
    version=settings.APP_VERSION,
    description="KVM 虚拟化管理平台 API",
)

# CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 注册路由
app.include_router(auth.router, prefix=f"{settings.API_PREFIX}/auth", tags=["认证"])
app.include_router(host_router.router, prefix=f"{settings.API_PREFIX}/hosts", tags=["主机管理"])
app.include_router(vm.router, prefix=f"{settings.API_PREFIX}/vm", tags=["虚拟机管理"])
app.include_router(storage.router, prefix=f"{settings.API_PREFIX}/storage", tags=["存储管理"])
app.include_router(network.router, prefix=f"{settings.API_PREFIX}/network", tags=["网络管理"])
app.include_router(snapshot.router, prefix=f"{settings.API_PREFIX}/snapshot", tags=["快照管理"])
app.include_router(monitor.router, prefix=f"{settings.API_PREFIX}/monitor", tags=["资源监控"])


@app.get("/")
async def root():
    return {"name": settings.APP_NAME, "version": settings.APP_VERSION}


@app.get(f"{settings.API_PREFIX}/host")
async def host_info():
    """获取本机宿主机信息（兼容旧接口）"""
    from app.libvirt_conn import libvirt_conn
    return libvirt_conn.get_host_info()


@app.get("/health")
async def health():
    """健康检查"""
    from app.libvirt_conn import conn_pool
    try:
        conn = conn_pool.get_conn("local")
        return {"status": "ok", "libvirt": conn.isAlive()}
    except Exception as e:
        return {"status": "error", "message": str(e)}


# VNC WebSocket 代理
from starlette.websockets import WebSocket
from starlette.responses import HTMLResponse
import asyncio
import socket
import struct


@app.websocket("/ws/vnc/{vm_name}")
async def vnc_websocket(websocket: WebSocket, vm_name: str, host_id: str = "local"):
    """WebSocket 代理到虚拟机 VNC"""
    await websocket.accept()

    from app.libvirt_conn import conn_pool
    from app.hosts import get_host as get_host_info
    from lxml import etree

    try:
        conn = conn_pool.get_conn(host_id)
        dom = conn.lookupByName(vm_name)
        xml_desc = dom.XMLDesc(0)
        tree = etree.fromstring(xml_desc.encode())

        # 获取 VNC 端口
        graphics = tree.find(".//graphics[@type='vnc']")
        if graphics is None:
            await websocket.close(code=1000, reason="虚拟机没有 VNC 配置")
            return

        vnc_port = int(graphics.get("port", -1))
        vnc_listen = graphics.get("listen", "127.0.0.1")

        if vnc_port <= 0:
            await websocket.close(code=1000, reason="VNC 端口未分配，虚拟机可能未运行")
            return

        # 根据主机类型决定 VNC 连接目标
        host_info = get_host_info(host_id)
        target_host = vnc_listen
        is_remote = host_info and host_info.type != "local"

        if is_remote:
            from urllib.parse import urlparse
            parsed = urlparse(host_info.uri)
            remote_host = parsed.hostname or "127.0.0.1"

            if vnc_listen in ("127.0.0.1", "localhost", ""):
                # VNC 只监听本地回环，需要 SSH 隧道
                if host_info.type == "ssh":
                    import subprocess, time
                    local_port = 20000 + (hash(vm_name + str(vnc_port)) % 10000)
                    # 先杀掉可能存在的旧隧道
                    subprocess.run(
                        ["fuser", "-k", f"{local_port}/tcp"],
                        capture_output=True, timeout=3,
                    )
                    ssh_args = ["ssh", "-f", "-N"]
                    if host_info.ssh_key_path:
                        ssh_args.extend(["-i", host_info.ssh_key_path])
                    ssh_args.extend([
                        "-o", "StrictHostKeyChecking=no",
                        "-o", "ServerAliveInterval=30",
                        "-L", f"127.0.0.1:{local_port}:127.0.0.1:{vnc_port}",
                        remote_host,
                    ])
                    result = subprocess.run(ssh_args, capture_output=True, timeout=10)
                    if result.returncode != 0:
                        err = result.stderr.decode(errors='replace').strip()
                        await websocket.close(code=1011, reason=f"SSH隧道建立失败: {err}")
                        return
                    target_host = "127.0.0.1"
                    vnc_port = local_port
                else:
                    # TCP 模式但 VNC 只听 localhost，无法直接连
                    await websocket.close(code=1011, reason="远程主机 VNC 监听在 localhost，需要使用 SSH 模式")
                    return
            elif vnc_listen == "0.0.0.0":
                # VNC 监听所有接口，直接连远程主机 IP
                target_host = remote_host
            else:
                # VNC 监听特定地址
                target_host = vnc_listen
        else:
            # 本地主机
            if vnc_listen in ("0.0.0.0", ""):
                target_host = "127.0.0.1"
            # else 保持 vnc_listen (127.0.0.1 或其他)

        # 连接到 VNC 服务器
        reader, writer = await asyncio.open_connection(target_host, vnc_port)

        async def ws_to_vnc():
            """WebSocket -> VNC"""
            try:
                while True:
                    data = await websocket.receive_bytes()
                    writer.write(data)
                    await writer.drain()
            except Exception:
                pass
            finally:
                writer.close()

        async def vnc_to_ws():
            """VNC -> WebSocket"""
            try:
                while True:
                    data = await reader.read(4096)
                    if not data:
                        break
                    await websocket.send_bytes(data)
            except Exception:
                pass
            finally:
                try:
                    await websocket.close()
                except Exception:
                    pass

        await asyncio.gather(ws_to_vnc(), vnc_to_ws())

    except Exception as e:
        try:
            await websocket.close(code=1011, reason=str(e))
        except Exception:
            pass