server.py

"""
Arista HUD Server
FastAPI backend with WebSocket push for real-time device telemetry.
Includes SSH terminal proxy via paramiko -> xterm.js.
"""

import asyncio
import json
import logging
import time
from contextlib import asynccontextmanager
from pathlib import Path

import paramiko
import yaml
import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles

from collector import AristaCollector

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("server")

# ---- Global state ----
CONFIG: dict = {}
collector: AristaCollector | None = None
clients: set[WebSocket] = set()
latest_data: dict | None = None
poll_task: asyncio.Task | None = None


def load_config() -> dict:
    cfg_path = Path(__file__).parent / "config.yaml"
    if not cfg_path.exists():
        raise FileNotFoundError(
            f"Config not found at {cfg_path}. Copy config.yaml and edit."
        )
    with open(cfg_path) as f:
        return yaml.safe_load(f)


async def poll_loop():
    """Background loop: collect data and push to all WebSocket clients."""
    global latest_data
    interval = CONFIG.get("poll_interval", 15)
    logger.info(f"Poll loop starting, interval={interval}s")

    while True:
        try:
            # Run blocking Netmiko call in thread pool
            loop = asyncio.get_event_loop()
            data = await loop.run_in_executor(None, collector.collect)
            latest_data = data

            # Broadcast to all connected clients
            payload = json.dumps(data, default=str)
            dead = set()
            for ws in clients:
                try:
                    await ws.send_text(payload)
                except Exception:
                    dead.add(ws)
            clients.difference_update(dead)

            if dead:
                logger.info(f"Removed {len(dead)} dead WebSocket client(s)")

        except Exception as e:
            logger.error(f"Poll loop error: {e}")

        await asyncio.sleep(interval)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Start the polling background task on startup."""
    global CONFIG, collector, poll_task

    CONFIG = load_config()
    collector = AristaCollector(CONFIG["device"])
    poll_task = asyncio.create_task(poll_loop())
    logger.info("Arista HUD server started")
    yield
    poll_task.cancel()
    logger.info("Arista HUD server stopped")


app = FastAPI(title="Arista HUD", lifespan=lifespan)

# Serve static files (the HUD frontend)
static_dir = Path(__file__).parent / "static"
if static_dir.exists():
    app.mount("/static", StaticFiles(directory=str(static_dir)), name="static")


@app.get("/")
async def root():
    """Serve the HUD frontend."""
    index = static_dir / "index.html"
    if index.exists():
        return FileResponse(str(index))
    return {"error": "static/index.html not found"}


@app.websocket("/ws")
async def websocket_endpoint(ws: WebSocket):
    """WebSocket endpoint for real-time telemetry push."""
    await ws.accept()
    clients.add(ws)
    logger.info(f"WebSocket client connected ({len(clients)} total)")

    # Send current data immediately if available
    if latest_data:
        try:
            await ws.send_text(json.dumps(latest_data, default=str))
        except Exception:
            pass

    try:
        # Keep connection alive, handle client messages (future: commands)
        while True:
            msg = await ws.receive_text()
            # Could handle client requests here (e.g., force refresh)
            if msg == "refresh":
                if latest_data:
                    await ws.send_text(json.dumps(latest_data, default=str))
    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.warning(f"WebSocket error: {e}")
    finally:
        clients.discard(ws)
        logger.info(f"WebSocket client disconnected ({len(clients)} remaining)")


@app.get("/api/data")
async def get_data():
    """REST fallback for current device state."""
    if latest_data:
        return latest_data
    return {"error": "No data collected yet"}


@app.get("/api/status")
async def get_status():
    """Server health check."""
    return {
        "status": "ok",
        "clients": len(clients),
        "poll_count": collector._collect_count if collector else 0,
        "last_error": collector._last_error if collector else None,
        "hostname": CONFIG.get("device", {}).get("host", "unknown"),
    }


def _open_ssh_shell(cfg: dict, cols: int = 120, rows: int = 36):
    """Open a paramiko interactive shell. Runs in thread pool."""
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    connect_kwargs = {
        "hostname": cfg["host"],
        "username": cfg["username"],
        "timeout": cfg.get("timeout", 30),
        "allow_agent": False,
        "look_for_keys": False,
    }

    key_file = cfg.get("key_file", "")
    if key_file:
        connect_kwargs["key_filename"] = str(Path(key_file).expanduser())
    if cfg.get("password"):
        connect_kwargs["password"] = cfg["password"]

    # Legacy SSH algorithm support
    if cfg.get("legacy_ssh", False):
        connect_kwargs["disabled_algorithms"] = {
            "pubkeys": ["rsa-sha2-256", "rsa-sha2-512"],
        }

    client.connect(**connect_kwargs)
    channel = client.invoke_shell(
        term="xterm-256color", width=cols, height=rows
    )
    channel.settimeout(0.05)
    return client, channel


@app.websocket("/ws/terminal")
async def terminal_ws(ws: WebSocket):
    """WebSocket SSH terminal proxy. Bridges xterm.js <-> paramiko shell."""
    await ws.accept()
    loop = asyncio.get_event_loop()
    device_cfg = CONFIG.get("device", {})
    client = None
    channel = None

    try:
        client, channel = await loop.run_in_executor(
            None, lambda: _open_ssh_shell(device_cfg)
        )
        logger.info(f"Terminal session opened to {device_cfg['host']}")
    except Exception as e:
        logger.error(f"Terminal SSH connect failed: {e}")
        await ws.send_text(f"\r\n\x1b[31mSSH connection failed: {e}\x1b[0m\r\n")
        await ws.close()
        return

    stop = asyncio.Event()

    async def ssh_reader():
        """Read from SSH channel, send to WebSocket."""
        while not stop.is_set():
            try:
                data = await loop.run_in_executor(
                    None, lambda: channel.recv(4096) if channel.recv_ready() else b""
                )
                if data:
                    await ws.send_bytes(data)
                elif channel.exit_status_ready():
                    break
                else:
                    await asyncio.sleep(0.02)
            except Exception:
                break
        stop.set()

    async def ws_reader():
        """Read from WebSocket, send to SSH channel."""
        try:
            while not stop.is_set():
                msg = await ws.receive()
                if msg.get("type") == "websocket.disconnect":
                    break
                data = msg.get("bytes") or (msg.get("text", "").encode())
                if data:
                    # Handle resize messages (JSON with cols/rows)
                    if data[:1] == b"{":
                        try:
                            resize = json.loads(data)
                            if "cols" in resize and "rows" in resize:
                                channel.resize_pty(
                                    width=resize["cols"], height=resize["rows"]
                                )
                                continue
                        except (json.JSONDecodeError, ValueError):
                            pass
                    await loop.run_in_executor(None, channel.send, data)
        except WebSocketDisconnect:
            pass
        except Exception as e:
            logger.warning(f"Terminal WS reader error: {e}")
        stop.set()

    try:
        await asyncio.gather(ssh_reader(), ws_reader())
    finally:
        if channel:
            try:
                channel.close()
            except Exception:
                pass
        if client:
            try:
                client.close()
            except Exception:
                pass
        logger.info("Terminal session closed")


if __name__ == "__main__":
    cfg = load_config()
    srv = cfg.get("server", {})
    uvicorn.run(
        "server:app",
        host=srv.get("host", "0.0.0.0"),
        port=srv.get("port", 8470),
        reload=False,
        log_level="info",
    )