zmb-webui/backend/main_aiohttp.py

"""
ZMB Webui API - aiohttp version (Low-RAM)
Single service: API + Static Files + WebSocket
"""

import asyncio
import json
import logging
import sys
from pathlib import Path
from typing import Set
from aiohttp import web
from aiohttp.web_runner import AppRunner, TCPSite

# Add backend to path
sys.path.insert(0, str(Path(__file__).parent))

from services.zfs_runner import zfs_runner
from services.auth import auth_service
from routers_aiohttp import setup_all_routes

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# WebSocket clients for broadcasting
ws_clients: Set[web.WebSocketResponse] = set()


async def ws_broadcast(data: dict):
    """Broadcast message to all connected WebSocket clients"""
    message = json.dumps(data)
    dead_clients = set()

    for ws in ws_clients:
        try:
            if not ws.is_closed():
                await ws.send_str(message)
            else:
                dead_clients.add(ws)
        except Exception as e:
            logger.error(f"WebSocket broadcast error: {e}")
            dead_clients.add(ws)

    # Clean up dead connections
    for ws in dead_clients:
        ws_clients.discard(ws)


async def pool_status_broadcaster():
    """Background task: broadcast pool status every 30 seconds"""
    while True:
        try:
            await asyncio.sleep(30)
            pools = zfs_runner.list_pools()
            if pools:
                await ws_broadcast({
                    "type": "pool_status",
                    "pools": pools
                })
        except Exception as e:
            logger.error(f"Pool status broadcast error: {e}")


async def startup(app):
    """Startup handler"""
    logger.info("Starting ZMB Webui API (aiohttp)")
    app['broadcaster_task'] = asyncio.create_task(pool_status_broadcaster())


async def shutdown(app):
    """Shutdown handler"""
    logger.info("Shutting down ZMB Webui API")
    if 'broadcaster_task' in app:
        app['broadcaster_task'].cancel()

    # Close all WebSocket connections
    for ws in ws_clients:
        await ws.close()


async def handle_websocket(request):
    """WebSocket endpoint for live pool updates"""
    ws = web.WebSocketResponse()
    await ws.prepare(request)
    ws_clients.add(ws)

    try:
        # Send initial pool status
        pools = zfs_runner.list_pools()
        await ws.send_json({
            "type": "initial",
            "pools": pools
        })

        # Keep connection open
        async for msg in ws:
            if msg.type == web.WSMsgType.TEXT:
                # Echo back (optional)
                pass
            elif msg.type == web.WSMsgType.ERROR:
                logger.error(f'WebSocket error: {ws.exception()}')
    finally:
        ws_clients.discard(ws)
        await ws.close()

    return ws


async def handle_status(request):
    """System status endpoint (no auth required)"""
    # Check if ZFS is available
    import subprocess
    try:
        result = subprocess.run(["zpool", "list"], capture_output=True, timeout=5)
        zfs_available = result.returncode == 0
    except Exception:
        zfs_available = False

    return web.json_response({
        "status": "healthy",
        "zfs_available": zfs_available,
        "version": "1.0.0"
    })


async def handle_health(request):
    """Health check endpoint"""
    return web.json_response({
        "status": "healthy",
        "version": "1.0.0"
    })


async def handle_static(request):
    """Serve static files (HTML, JS, CSS)"""
    path = request.match_info['path']

    # Security: prevent directory traversal
    if '..' in path or path.startswith('/'):
        return web.Response(status=400, text="Invalid path")

    # Try to find file in frontend/out directory
    static_dir = Path(__file__).parent.parent / "frontend" / "out"
    file_path = static_dir / path

    # Ensure file is within static_dir
    try:
        file_path.resolve().relative_to(static_dir.resolve())
    except ValueError:
        return web.Response(status=403, text="Forbidden")

    if file_path.exists() and file_path.is_file():
        return web.FileResponse(file_path)

    # If requesting a directory or file not found, try index.html (SPA routing)
    index_path = static_dir / "index.html"
    if index_path.exists():
        return web.FileResponse(index_path)

    return web.Response(status=404, text="Not found")


async def handle_root(request):
    """Root endpoint - serve index.html"""
    static_dir = Path(__file__).parent.parent / "frontend" / "out"
    index_path = static_dir / "index.html"

    if index_path.exists():
        return web.FileResponse(index_path)

    return web.json_response({
        "name": "ZMB Webui API",
        "version": "1.0.0",
        "docs": "/api/docs",
        "login": "/login"
    })


def create_app():
    """Create and configure the aiohttp application"""
    app = web.Application()

    # Store WebSocket clients and broadcaster task
    app['ws_clients'] = ws_clients

    # Startup/Shutdown handlers
    app.on_startup.append(startup)
    app.on_shutdown.append(shutdown)

    # Routes
    app.router.add_get('/health', handle_health)
    app.router.add_get('/api/status', handle_status)
    app.router.add_get('/ws', handle_websocket)

    # API Routes (all routers)
    setup_all_routes(app)

    # Static file serving (must be last - catch-all)
    app.router.add_get('/{path_info:.*}', handle_static)
    app.router.add_get('/', handle_root)

    return app


async def main():
    """Main entry point"""
    app = create_app()
    runner = AppRunner(app)
    await runner.setup()

    # Start server on port 8000
    site = TCPSite(runner, '0.0.0.0', 8000)
    await site.start()

    logger.info("ZMB Webui API listening on http://0.0.0.0:8000")
    logger.info("Available endpoints:")
    logger.info("  GET  /health              - Health check")
    logger.info("  GET  /api/status          - System status (ZFS available)")
    logger.info("  POST /api/auth/login      - Login with username/password")
    logger.info("  GET  /api/pools           - List all pools")
    logger.info("  GET  /api/datasets        - List datasets")
    logger.info("  GET  /api/snapshots       - List snapshots")
    logger.info("  GET  /ws                  - WebSocket (live updates)")
    logger.info("  GET  /* (except /api)     - Static files (HTML/CSS/JS)")
    logger.info("")

    # Keep running
    try:
        await asyncio.Event().wait()
    except KeyboardInterrupt:
        logger.info("Shutting down...")
        await runner.cleanup()


if __name__ == "__main__":
    asyncio.run(main())