pi-pm3/app/backend/main.py

"""Main FastAPI application for Dangerous Pi."""
import asyncio
import os
from pathlib import Path
from contextlib import asynccontextmanager
from fastapi import FastAPI, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles

from . import config
from .models.database import init_db
from .api import health, pm3, system, wifi, updates, plugins
from .api.auth import verify_credentials
from .websocket import router as ws_router
from .websocket import notifications as ws_notifications
from .managers.update_manager import get_update_manager
from .managers.ups_manager import get_ups_manager
from .managers.ble_manager import get_ble_manager, NotificationType
from .managers.plugin_manager import get_plugin_manager
from .managers.wifi_manager import wifi_manager
from .services.container import container


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan manager."""
    # Startup
    print(f"🚀 Starting Dangerous Pi backend...")
    await init_db()
    print(f"✅ Database initialized")

    # Start periodic update checks
    update_manager = get_update_manager()
    update_task = asyncio.create_task(update_manager.start_periodic_checks())
    print(f"✅ Update manager started")

    # Initialize and start BLE manager
    ble_manager = get_ble_manager()
    await ble_manager.initialize()
    if ble_manager.is_available():
        await ble_manager.start_advertising()
        print(f"✅ BLE manager started")
    else:
        print(f"⚠️  BLE not available")

    # Apply saved WiFi mode (for boot persistence)
    try:
        await wifi_manager.apply_saved_mode()
        print(f"✅ WiFi mode restored")
    except Exception as e:
        print(f"⚠️  WiFi mode restore failed: {e}")

    # Start UPS monitoring
    ups_manager = get_ups_manager()

    # Register UPS event callbacks for WebSocket notifications and BLE
    async def ups_event_handler(event):
        """Handle UPS events and broadcast via WebSocket and BLE."""
        event_type = event.get("type")
        data = event.get("data", {})

        if event_type == "battery_warning":
            await ws_notifications.notify_ups_warning(data["percentage"], data["threshold"])
            await ble_manager.send_notification(
                NotificationType.BATTERY_WARNING,
                f"Battery low: {data['percentage']:.1f}%",
                data
            )
        elif event_type == "battery_low":
            await ws_notifications.notify_ups_critical(data["percentage"])
            await ble_manager.send_notification(
                NotificationType.BATTERY_CRITICAL,
                f"Battery critical: {data['percentage']:.1f}%",
                data
            )
        elif event_type == "battery_critical":
            await ws_notifications.notify_ups_shutdown(data.get("delay", 60), data["percentage"])
            await ble_manager.send_notification(
                NotificationType.SHUTDOWN_INITIATED,
                f"Shutdown initiated: {data['percentage']:.1f}% battery",
                data
            )
        elif event_type == "ups_config_required":
            await ws_notifications.notify_ups_config_required(
                data.get("model", "Unknown"),
                data.get("message", "UPS configuration required"),
                data.get("hint", "")
            )
            await ble_manager.send_notification(
                NotificationType.CONFIG_REQUIRED,
                data.get("message", "UPS configuration required"),
                data
            )

    ups_manager.register_event_callback(ups_event_handler)
    ups_task = asyncio.create_task(ups_manager.start_monitoring())
    print(f"✅ UPS manager started")

    # Discover and load plugins
    plugin_manager = get_plugin_manager()
    discovered = await plugin_manager.discover_plugins()
    print(f"✅ Plugin manager started ({len(discovered)} plugins discovered)")

    # Start PM3 device manager for multi-device support
    pm3_device_manager = container.pm3_device_manager

    # Register PM3 device change callback
    async def pm3_device_change_handler(device_list):
        """Handle PM3 device list changes and broadcast via WebSocket."""
        devices_data = [
            {
                "device_id": d.device_id,
                "device_path": d.device_path,
                "status": d.status.value if hasattr(d.status, 'value') else d.status,
                "friendly_name": d.friendly_name,
            }
            for d in device_list
        ]
        await ws_notifications.notify_pm3_devices(devices_data)

    pm3_device_manager.register_device_change_callback(pm3_device_change_handler)
    await pm3_device_manager.start()
    print(f"✅ PM3 device manager started")

    # Start periodic system stats broadcasting (every 5 seconds)
    async def broadcast_system_stats():
        """Periodically broadcast system stats via WebSocket."""
        while True:
            try:
                result = await container.system_service.get_info()
                if result.success:
                    cpu_data = result.data["cpu"]
                    memory_data = result.data["memory"]
                    await ws_notifications.notify_system_stats(
                        cpu_percent=cpu_data.get("percent", 0),
                        cpu_per_core=cpu_data.get("per_core", []),
                        load_average=cpu_data.get("load_average", []),
                        memory_percent=memory_data.get("percent", 0),
                        temperature=cpu_data.get("temperature")
                    )
            except Exception as e:
                print(f"Error broadcasting system stats: {e}")
            await asyncio.sleep(5)  # Update every 5 seconds

    system_stats_task = asyncio.create_task(broadcast_system_stats())
    print(f"✅ System stats broadcaster started")

    yield

    # Shutdown
    print(f"🛑 Shutting down Dangerous Pi backend...")

    # Stop PM3 device manager
    await pm3_device_manager.stop()
    print(f"✅ PM3 device manager stopped")

    update_task.cancel()
    ups_task.cancel()
    system_stats_task.cancel()
    try:
        await update_task
    except asyncio.CancelledError:
        pass
    try:
        await ups_task
    except asyncio.CancelledError:
        pass
    try:
        await system_stats_task
    except asyncio.CancelledError:
        pass

    # Close UPS manager I2C connection
    ups_manager.close()


app = FastAPI(
    title="Dangerous Pi API",
    description="Backend API for Proxmark3 management on Raspberry Pi Zero 2 W",
    version="0.1.0",
    lifespan=lifespan
)

# CORS middleware for frontend
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Permissive for development; restrict via CORS_ORIGINS env var in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Auth dependency for protected routes (applied when AUTH_ENABLED=true)
auth_dependency = [Depends(verify_credentials)] if config.AUTH_ENABLED else []

# Include routers - all protected when AUTH_ENABLED=true
app.include_router(health.router, prefix="/api", tags=["health"], dependencies=auth_dependency)
app.include_router(pm3.router, prefix="/api/pm3", tags=["proxmark3"], dependencies=auth_dependency)
app.include_router(system.router, prefix="/api/system", tags=["system"], dependencies=auth_dependency)
app.include_router(wifi.router, prefix="/api/wifi", tags=["wifi"], dependencies=auth_dependency)
app.include_router(updates.router, prefix="/api/updates", tags=["updates"], dependencies=auth_dependency)
app.include_router(plugins.router, prefix="/api/plugins", tags=["plugins"], dependencies=auth_dependency)
app.include_router(ws_router, prefix="/ws", tags=["websocket"])

# Serve frontend static files if build directory exists
# In production, frontend is built and served from /opt/dangerous-pi/app/frontend/build
# Priority: 1) check relative to working directory, 2) check absolute production path
frontend_build_paths = [
    Path(__file__).parent.parent / "frontend" / "build" / "client",  # Development
    Path("/opt/dangerous-pi/app/frontend/build/client"),              # Production
]

frontend_path = None
for path in frontend_build_paths:
    if path.exists() and path.is_dir():
        frontend_path = path
        break

if frontend_path:
    # Mount static assets (JS, CSS, images)
    assets_path = frontend_path / "assets"
    if assets_path.exists():
        app.mount("/assets", StaticFiles(directory=str(assets_path)), name="assets")

    # Cache control headers
    NO_CACHE_HEADERS = {
        "Cache-Control": "no-cache, no-store, must-revalidate",
        "Pragma": "no-cache",
        "Expires": "0"
    }

    # Serve index.html for all non-API routes (SPA support)
    @app.get("/{full_path:path}")
    async def serve_frontend(full_path: str):
        """Serve frontend files, fall back to index.html for SPA routing."""
        # Check for exact file match first
        file_path = frontend_path / full_path
        if file_path.exists() and file_path.is_file():
            # HTML files get no-cache headers
            if str(file_path).endswith('.html'):
                return FileResponse(str(file_path), headers=NO_CACHE_HEADERS)
            return FileResponse(str(file_path))
        # Fall back to index.html for SPA routing (no-cache for HTML)
        index_path = frontend_path / "index.html"
        if index_path.exists():
            return FileResponse(str(index_path), headers=NO_CACHE_HEADERS)
        return JSONResponse(status_code=404, content={"error": "Not found"})

    print(f"📁 Serving frontend from: {frontend_path}")
else:
    print(f"⚠️  Frontend build not found, API-only mode")


@app.exception_handler(Exception)
async def global_exception_handler(request, exc):
    """Global exception handler."""
    return JSONResponse(
        status_code=500,
        content={"error": str(exc)}
    )


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        "app.backend.main:app",
        host=config.HOST,
        port=config.PORT,
        reload=True
    )