qemu/web_interface.py

"""
Web Interface for QEMU VM
Provides browser-based access to VM display
"""

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse
import asyncio
import logging
from pathlib import Path
from typing import Optional, Dict
import json
import os
import sys
import subprocess

# Add parent directory to path for imports
current_dir = Path(__file__).parent
root_dir = current_dir.parent
if str(root_dir) not in sys.path:
    sys.path.insert(0, str(root_dir))

from qemu_manager import QEMUManager

# Configure logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

app = FastAPI()

# Mount static files
static_path = Path(__file__).parent / "static"
app.mount("/static", StaticFiles(directory=str(static_path)), name="static")

# Global QEMU manager instance
qemu: Optional[QEMUManager] = None
websockify_process: Optional[subprocess.Popen] = None

class VMConnection:

    def __init__(self):
        self.active_connections: Dict[int, WebSocket] = {}
        
    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        client_id = id(websocket)
        self.active_connections[client_id] = websocket
        return client_id
        
    def disconnect(self, client_id: int):
        self.active_connections.pop(client_id, None)
        
    async def broadcast(self, message: str):
        for ws in self.active_connections.values():
            try:
                await ws.send_text(message)
            except:
                continue

vm_connection = VMConnection()

def start_websockify(vnc_port: int, websocket_port: int = 6080):
    """Start websockify proxy for VNC"""
    global websockify_process
    
    if websockify_process:
        websockify_process.terminate()
        websockify_process.wait()
    
    cmd = [
        "websockify",
        f"{websocket_port}",
        f"localhost:{vnc_port}"
    ]
    
    websockify_process = subprocess.Popen(cmd)
    logger.info(f"Started websockify proxy on port {websocket_port} for VNC port {vnc_port}")
    return websocket_port

@app.get("/")
async def get_index():
    """Serve main HTML interface"""
    html_path = static_path / "index.html"
    return HTMLResponse(content=html_path.read_text())

@app.websocket("/ws/vm")
async def vm_websocket(websocket: WebSocket):
    """Handle VM display and control WebSocket"""
    client_id = await vm_connection.connect(websocket)
    
    try:
        while True:
            message = await websocket.receive_json()
            
            if message["type"] == "install":
                # Start fresh OS installation
                if qemu:
                    vnc_port = await qemu.install_os(message["iso_url"])
                    websocket_port = start_websockify(vnc_port)
                    await websocket.send_json({
                        "type": "vnc_info",
                        "port": websocket_port
                    })
                    
            elif message["type"] == "boot":
                # Boot existing OS
                if qemu:
                    try:
                        vnc_port = await qemu.boot_os()
                        websocket_port = start_websockify(vnc_port)
                        await websocket.send_json({
                            "type": "vnc_info",
                            "port": websocket_port
                        })
                    except FileNotFoundError:
                        await websocket.send_json({
                            "type": "error",
                            "message": "No OS installation found"
                        })
                        
            elif message["type"] == "shutdown":
                # Shutdown VM
                if qemu:
                    await qemu.shutdown()
                    await websocket.send_json({
                        "type": "status",
                        "message": "VM shut down"
                    })
                    
    except WebSocketDisconnect:
        vm_connection.disconnect(client_id)
        
@app.on_event("startup")
async def startup_event():
    """Initialize QEMU manager on startup"""
    global qemu
    qemu = QEMUManager()
    logger.info("QEMU manager initialized")
    
@app.on_event("shutdown")
async def shutdown_event():
    """Cleanup on shutdown"""
    global websockify_process
    
    if qemu:
        await qemu.shutdown()
        logger.info("QEMU manager shut down")
        
    if websockify_process:
        websockify_process.terminate()
        websockify_process.wait()
        logger.info("Websockify proxy shut down")


def main():
    """Entry point for the FServe application"""
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8080)

if __name__ == "__main__":
    main()