stackpilot/backend/routers/ws.py

"""WebSocket endpoints for real-time log streaming and Docker events."""
from __future__ import annotations

import asyncio
import json
import logging

import contextlib

from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect
from jose import JWTError
from sqlmodel import Session

from auth import decode_token, resolve_token_user
from database import engine
from models.setting import EVENT_PULL_FAILED, EVENT_STACK_ERROR, EVENT_STACK_START
from services import (
    audit_service,
    compose_service,
    exec_service,
    notify_service,
    stack_lock_service,
    update_service,
)

logger = logging.getLogger("stackpilot.ws")

router = APIRouter(tags=["ws"])


def _user_for(token: str | None):
    """The live user behind a socket's token, or None.

    Resolves against the database rather than reading the role straight off the
    JWT: a socket can outlive a demotion, a disabled account or a password
    reset, and the exec endpoint below is root-equivalent on the host. Same
    check the HTTP routes make.
    """
    if not token:
        return None
    try:
        payload = decode_token(token, "access")
    except (JWTError, Exception):  # noqa: BLE001
        return None
    with Session(engine) as session:
        user = resolve_token_user(session, payload)
        if user:
            session.expunge(user)
        return user


async def _authorize(websocket: WebSocket, token: str | None) -> bool:
    """Validate the JWT supplied as a query param. Closes socket on failure."""
    if _user_for(token) is None:
        await websocket.close(code=4401)
        return False
    return True


async def _authorize_admin(websocket: WebSocket, token: str | None) -> bool:
    """Like _authorize but also requires the admin role (exec is root-equivalent).

    Closes 4401 on a missing/invalid token, 4403 on a valid non-admin token."""
    user = _user_for(token)
    if user is None:
        await websocket.close(code=4401)
        return False
    if user.role != "admin":
        await websocket.close(code=4403)
        return False
    return True


async def _stream_logs(websocket: WebSocket, stack_id: str, service: str | None):
    """Stream `docker compose logs -f` output to the client."""
    args = ["logs", "--no-color", "--tail", "200", "--timestamps", "-f"]
    if service:
        args.append(service)
    try:
        async for line in compose_service.stream_compose(stack_id, args):
            await websocket.send_text(
                json.dumps(
                    {
                        "type": "log",
                        "stack_id": stack_id,
                        "service": service,
                        "line": line,
                    }
                )
            )
    except WebSocketDisconnect:
        raise
    except Exception as exc:  # noqa: BLE001
        await websocket.send_text(
            json.dumps({"type": "error", "detail": str(exc)})
        )


@router.websocket("/ws/logs/{stack_id}")
async def ws_stack_logs(
    websocket: WebSocket,
    stack_id: str,
    token: str | None = Query(default=None),
):
    await websocket.accept()
    if not await _authorize(websocket, token):
        return
    try:
        await _stream_logs(websocket, stack_id, None)
    except WebSocketDisconnect:
        pass


@router.websocket("/ws/logs/{stack_id}/{service}")
async def ws_service_logs(
    websocket: WebSocket,
    stack_id: str,
    service: str,
    token: str | None = Query(default=None),
):
    await websocket.accept()
    if not await _authorize(websocket, token):
        return
    try:
        await _stream_logs(websocket, stack_id, service)
    except WebSocketDisconnect:
        pass


@router.websocket("/ws/deploy/{stack_id}")
async def ws_deploy(
    websocket: WebSocket,
    stack_id: str,
    token: str | None = Query(default=None),
):
    """Run `docker compose up -d` and stream its output (image pulls, container
    creation) to the browser so the user sees deploy progress live. Records the
    same audit entry and notification as the REST `/start` endpoint."""
    await websocket.accept()
    if not await _authorize(websocket, token):
        return
    username = decode_token(token, "access").get("sub", "unknown") if token else "unknown"

    rc: int | None = None
    disconnected = False
    # Same guard the REST lifecycle uses — the deploy console runs the very
    # same `compose up`, so it has to queue behind an in-flight operation
    # rather than race it.
    lock_session = Session(engine)
    try:
        stack_lock_service.acquire(lock_session, stack_id, "start", username)
    except stack_lock_service.StackBusy as exc:
        lock_session.close()
        with contextlib.suppress(Exception):
            await websocket.send_text(json.dumps({"type": "error", "detail": str(exc)}))
            await websocket.close(code=4409)
        return
    try:
        async for kind, payload in compose_service.stream_up(stack_id):
            if kind == "log":
                await websocket.send_text(json.dumps({"type": "log", "line": payload}))
            else:
                rc = payload
                await websocket.send_text(json.dumps({"type": "done", "returncode": rc}))
    except WebSocketDisconnect:
        # Client navigated away; the compose subprocess keeps running so the
        # deploy still completes in the background.
        disconnected = True
    except Exception as exc:  # noqa: BLE001
        with contextlib.suppress(Exception):
            await websocket.send_text(json.dumps({"type": "error", "detail": str(exc)}))
    finally:
        with contextlib.suppress(Exception):
            stack_lock_service.release(lock_session, stack_id)
        lock_session.close()

    ok = rc in (0, None)
    try:
        with Session(engine) as session:
            audit_service.record(
                session, user=username, action="stack.start", target=stack_id,
                detail=f"rc={rc} (deploy console)", ip="ws",
            )
            if ok:
                await notify_service.notify(
                    EVENT_STACK_START, f"Stack '{stack_id}' started",
                    "compose up completed successfully.", session,
                )
            else:
                await notify_service.notify(
                    EVENT_STACK_ERROR, f"Stack '{stack_id}' start failed",
                    "compose up returned a non-zero exit code.", session,
                )
    except Exception:  # noqa: BLE001 - audit/notify are best-effort
        pass
    if not disconnected:
        with contextlib.suppress(Exception):
            await websocket.close()


@router.websocket("/ws/update/{stack_id}")
async def ws_update(
    websocket: WebSocket,
    stack_id: str,
    token: str | None = Query(default=None),
):
    """Run `docker compose pull && up -d` and stream its output, so the stacks
    list can render real update progress. Same audit/notify contract as the
    REST `/update` endpoint, which stays for non-interactive callers."""
    await websocket.accept()
    if not await _authorize_admin(websocket, token):
        return
    username = decode_token(token, "access").get("sub", "unknown") if token else "unknown"

    rc: int | None = None
    disconnected = False
    lock_session = Session(engine)
    try:
        stack_lock_service.acquire(lock_session, stack_id, "update", username)
    except stack_lock_service.StackBusy as exc:
        lock_session.close()
        with contextlib.suppress(Exception):
            await websocket.send_text(json.dumps({"type": "error", "detail": str(exc)}))
            await websocket.close(code=4409)
        return
    try:
        async for kind, payload in compose_service.stream_update(stack_id):
            if kind == "log":
                await websocket.send_text(json.dumps({"type": "log", "line": payload}))
            else:
                rc = payload
                await websocket.send_text(json.dumps({"type": "done", "returncode": rc}))
    except WebSocketDisconnect:
        # Client navigated away; compose keeps running so the update finishes.
        disconnected = True
    except Exception as exc:  # noqa: BLE001
        with contextlib.suppress(Exception):
            await websocket.send_text(json.dumps({"type": "error", "detail": str(exc)}))
    finally:
        with contextlib.suppress(Exception):
            stack_lock_service.release(lock_session, stack_id)
        lock_session.close()

    ok = rc in (0, None)
    try:
        with Session(engine) as session:
            audit_service.record(
                session, user=username, action="stack.update", target=stack_id,
                detail=f"rc={rc} (update stream)", ip="ws",
            )
            if ok:
                await notify_service.notify(
                    EVENT_STACK_START, f"Stack '{stack_id}' updated",
                    "compose pull + up completed successfully.", session,
                )
            else:
                await notify_service.notify(
                    EVENT_PULL_FAILED, f"Stack '{stack_id}' update failed",
                    "compose pull/up returned a non-zero exit code.", session,
                )
    except Exception:  # noqa: BLE001 - audit/notify are best-effort
        pass
    if not disconnected:
        with contextlib.suppress(Exception):
            await websocket.close()

    if ok:
        update_service.refresh_stack_local(stack_id)


#: Docker event types worth telling the UI about. Filtered daemon-side, so the
#: bulk of the firehose never crosses the socket.
_EVENT_TYPES = ["container", "image", "network", "volume"]

#: Container actions that say nothing about state a page renders. exec_* alone
#: is three events per web-terminal keystroke session, and `top`/`attach` fire
#: whenever something inspects a container — invalidating queries on those would
#: make the stream noisier than the polling it replaces.
_IGNORED_ACTIONS = {
    "exec_create", "exec_start", "exec_die", "exec_detach",
    "attach", "top", "resize", "archive-path", "extract-to-dir",
}


@router.websocket("/ws/events")
async def ws_events(
    websocket: WebSocket,
    token: str | None = Query(default=None),
):
    """Stream Docker events so the UI can refresh on change instead of polling.

    Every page used to poll its own endpoint every few seconds. Almost all of
    that state only changes when Docker does something, which is exactly what
    this reports — so the client refreshes on an event and keeps a slow poll as
    a safety net.

    Payload per event::

        {"type": "event", "resource": "container", "action": "start",
         "container": "jellyfin", "stack": "jellyfin"}

    ``resource`` is what the client needs to decide which queries to drop.
    """
    await websocket.accept()
    if not await _authorize(websocket, token):
        return
    from docker_client import get_client

    loop = asyncio.get_event_loop()
    queue: asyncio.Queue = asyncio.Queue()
    stream = None

    def reader():
        """Blocking read of the event stream, handed to the loop thread-safely."""
        nonlocal stream
        try:
            stream = get_client().events(decode=True, filters={"type": _EVENT_TYPES})
            for event in stream:
                loop.call_soon_threadsafe(queue.put_nowait, event)
        except Exception:  # noqa: BLE001 - a closed stream lands here on teardown
            pass
        finally:
            loop.call_soon_threadsafe(queue.put_nowait, None)

    task = loop.run_in_executor(None, reader)
    try:
        await websocket.send_text(json.dumps({"type": "ready"}))
        while True:
            event = await queue.get()
            if event is None:  # reader finished — daemon gone or stream closed
                await websocket.send_text(
                    json.dumps({"type": "error", "detail": "Docker event stream ended"})
                )
                break
            action = (event.get("Action") or "").split(":")[0]
            if action in _IGNORED_ACTIONS:
                continue
            actor = event.get("Actor", {}) or {}
            attrs = actor.get("Attributes", {}) or {}
            await websocket.send_text(
                json.dumps(
                    {
                        "type": "event",
                        "resource": event.get("Type"),
                        "action": event.get("Action"),
                        "container": attrs.get("name"),
                        "stack": attrs.get("com.docker.compose.project"),
                    }
                )
            )
    except WebSocketDisconnect:
        pass
    except Exception as exc:  # noqa: BLE001
        with contextlib.suppress(Exception):
            await websocket.send_text(json.dumps({"type": "error", "detail": str(exc)}))
    finally:
        # Closing the stream is what actually unblocks the reader thread.
        # Cancelling the executor future does not: a thread already inside a
        # blocking read keeps that read, and the thread leaks for the life of
        # the process — once per page load, with the socket held open.
        if stream is not None:
            with contextlib.suppress(Exception):
                stream.close()
        task.cancel()


@router.websocket("/ws/exec/{container_id}")
async def ws_exec(
    websocket: WebSocket,
    container_id: str,
    token: str | None = Query(default=None),
    cmd: str | None = Query(default=None),
):
    """Interactive shell into a compose-managed container (admin only)."""
    await websocket.accept()
    if not await _authorize_admin(websocket, token):
        return
    username = decode_token(token, "access").get("sub", "unknown") if token else "unknown"
    shell = cmd or exec_service.DEFAULT_SHELL

    try:
        exec_id = exec_service.create_exec(container_id, [shell])
        holder, raw = exec_service.start_exec(exec_id)
    except Exception as exc:  # noqa: BLE001
        with contextlib.suppress(Exception):
            await websocket.send_text(json.dumps({"type": "error", "detail": str(exc)}))
        with contextlib.suppress(Exception):
            await websocket.close()
        return

    with contextlib.suppress(Exception):
        with Session(engine) as session:
            audit_service.record(
                session, user=username, action="container.exec",
                target=container_id[:12], detail=shell, ip="ws",
            )

    try:
        await exec_service.pump_exec(websocket, exec_id, holder, raw)
    except WebSocketDisconnect:
        pass
    finally:
        with contextlib.suppress(Exception):
            await websocket.close()