stackpilot/backend/agent_app.py

"""StackPilot agent — a slim, token-guarded Docker Compose API for one host.

The agent runs on each remote host (same image as the backend, different CMD).
It has no users, no database and no UI: it exposes just enough of the stack /
system surface for a central StackPilot to manage this host's compose stacks,
authenticated by a single shared bearer token (``AGENT_TOKEN``).

All compose/Docker logic is reused from the backend's ``compose_service`` and
``docker_client`` so behaviour matches the local host exactly.
"""
from __future__ import annotations

import logging
import os
from dataclasses import asdict

import tempfile

from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, Query, Request, UploadFile
from fastapi.responses import FileResponse, JSONResponse
from pydantic import BaseModel

from config import settings
from docker_client import DockerError, get_client, safe_call
from services import backup_service, compose_service

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

AGENT_VERSION = "0.8.0"


# --------------------------------------------------------------------------- #
# Auth
# --------------------------------------------------------------------------- #


def verify_token(authorization: str = Header(default="")) -> None:
    expected = settings.AGENT_TOKEN
    if not expected:
        raise HTTPException(status_code=503, detail="Agent token not configured")
    if authorization != f"Bearer {expected}":
        raise HTTPException(status_code=401, detail="Invalid agent token")


# --------------------------------------------------------------------------- #
# Schemas
# --------------------------------------------------------------------------- #


class StackBody(BaseModel):
    name: str | None = None
    yaml: str | None = None
    env: str | None = None


# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #


def _summary(stack_id: str) -> dict:
    try:
        containers = compose_service.containers_for_stack(stack_id)
        status = compose_service.compute_status(stack_id)
    except DockerError:
        containers = []
        status = "unknown"
    return {
        "id": stack_id,
        "name": stack_id,
        "description": None,
        "status": status,
        "service_count": len(containers),
        "running_count": sum(1 for c in containers if c.state == "running"),
        "created_at": None,
        "updated_at": None,
    }


def _hostname() -> str:
    return os.uname().nodename


def _system_info() -> dict:
    docker_version = ""
    host_os = ""
    running = total = 0
    try:
        client = get_client()
        docker_version = safe_call(client.version).get("Version", "")
        info = safe_call(client.info)
        host_os = info.get("OperatingSystem", "")
        running = info.get("ContainersRunning", 0)
        total = info.get("Containers", 0)
    except DockerError as exc:
        docker_version = f"unavailable ({exc.error})"
    return {
        "hostname": _hostname(),
        "docker_version": docker_version,
        "host_os": host_os,
        "containers_running": running,
        "containers_total": total,
    }


# --------------------------------------------------------------------------- #
# App
# --------------------------------------------------------------------------- #

app = FastAPI(title="StackPilot Agent", version=AGENT_VERSION)


@app.exception_handler(DockerError)
async def _docker_error(_request: Request, exc: DockerError):
    return JSONResponse(status_code=502, content={"error": exc.error, "detail": exc.detail})


@app.get("/agent/ping", dependencies=[Depends(verify_token)])
def ping() -> dict:
    return {"ok": True, "hostname": _hostname(), "version": AGENT_VERSION}


@app.get("/agent/system", dependencies=[Depends(verify_token)])
def system() -> dict:
    return _system_info()


@app.get("/agent/stacks", dependencies=[Depends(verify_token)])
def list_stacks() -> list[dict]:
    return [_summary(sid) for sid in compose_service.discover_stacks()]


@app.get("/agent/stacks/{stack_id}", dependencies=[Depends(verify_token)])
def get_stack(stack_id: str) -> dict:
    directory = compose_service.stack_dir(stack_id)
    if not os.path.isdir(directory):
        raise HTTPException(status_code=404, detail=f"Stack '{stack_id}' not found")
    try:
        containers = [asdict(c) for c in compose_service.containers_for_stack(stack_id)]
        status = compose_service.compute_status(stack_id)
    except DockerError:
        containers = []
        status = "unknown"
    return {
        "id": stack_id,
        "name": stack_id,
        "description": None,
        "status": status,
        "yaml": compose_service.read_compose(stack_id),
        "env": compose_service.read_env(stack_id),
        "containers": containers,
        "created_at": None,
        "updated_at": None,
    }


@app.post("/agent/stacks", dependencies=[Depends(verify_token)], status_code=201)
def create_stack(body: StackBody) -> dict:
    if not body.name:
        raise HTTPException(status_code=400, detail="name is required")
    stack_id = compose_service.slugify(body.name)
    if os.path.isdir(compose_service.stack_dir(stack_id)):
        raise HTTPException(status_code=409, detail=f"Stack '{stack_id}' already exists")
    compose_service.write_compose(stack_id, body.yaml or "services:\n")
    if body.env:
        compose_service.write_env(stack_id, body.env)
    return _summary(stack_id)


@app.put("/agent/stacks/{stack_id}", dependencies=[Depends(verify_token)])
def update_stack(stack_id: str, body: StackBody) -> dict:
    if not os.path.isdir(compose_service.stack_dir(stack_id)):
        raise HTTPException(status_code=404, detail=f"Stack '{stack_id}' not found")
    if body.yaml is not None:
        compose_service.write_compose(stack_id, body.yaml)
    if body.env is not None:
        compose_service.write_env(stack_id, body.env)
    return _summary(stack_id)


@app.delete("/agent/stacks/{stack_id}", dependencies=[Depends(verify_token)])
async def delete_stack(stack_id: str, delete_files: bool = Query(True)) -> dict:
    if not os.path.isdir(compose_service.stack_dir(stack_id)):
        raise HTTPException(status_code=404, detail=f"Stack '{stack_id}' not found")
    try:
        await compose_service.down(stack_id)
    except Exception:  # noqa: BLE001 - best-effort teardown
        pass
    if delete_files:
        compose_service.delete_stack_files(stack_id)
    return {"ok": True}


_ACTIONS = {
    "start": compose_service.up,
    "stop": compose_service.stop,
    "restart": compose_service.restart,
    "pull": compose_service.pull,
    "update": compose_service.update,
    "down": compose_service.down,
}


@app.post("/agent/stacks/{stack_id}/{action}", dependencies=[Depends(verify_token)])
async def lifecycle(stack_id: str, action: str) -> dict:
    fn = _ACTIONS.get(action)
    if not fn:
        raise HTTPException(status_code=400, detail=f"Unknown action '{action}'")
    if not os.path.isdir(compose_service.stack_dir(stack_id)):
        raise HTTPException(status_code=404, detail=f"Stack '{stack_id}' not found")
    result = await fn(stack_id)
    if result.get("returncode") not in (0, None):
        raise HTTPException(
            status_code=500,
            detail={
                "error": f"compose {action} failed",
                "detail": result.get("stderr", "").strip()[-2000:],
            },
        )
    return result


@app.get("/agent/stacks/{stack_id}/logs", dependencies=[Depends(verify_token)])
async def stack_logs(stack_id: str, tail: int = Query(200, le=2000)) -> dict:
    if not os.path.isdir(compose_service.stack_dir(stack_id)):
        raise HTTPException(status_code=404, detail=f"Stack '{stack_id}' not found")
    result = await compose_service.logs(stack_id, tail=tail)
    return {"logs": result.get("stdout", "") + result.get("stderr", "")}


@app.get("/agent/stacks/{stack_id}/backup", dependencies=[Depends(verify_token)])
async def backup_stack(
    stack_id: str,
    include_volumes: bool = Query(True),
    stop_first: bool = Query(True),
):
    if not os.path.isdir(compose_service.stack_dir(stack_id)):
        raise HTTPException(status_code=404, detail=f"Stack '{stack_id}' not found")
    try:
        path = await backup_service.create_backup(
            stack_id, stack_id, include_volumes=include_volumes, stop_first=stop_first,
        )
    except backup_service.BackupError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return FileResponse(
        path,
        media_type="application/gzip",
        filename=backup_service.backup_filename(stack_id, include_volumes),
    )


@app.post("/agent/stacks/restore", dependencies=[Depends(verify_token)])
async def restore_stack(
    file: UploadFile = File(...),
    target_id: str | None = Form(None),
    overwrite: bool = Form(False),
    restore_volumes: bool = Form(True),
) -> dict:
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".tar.gz")
    try:
        while chunk := await file.read(1024 * 1024):
            tmp.write(chunk)
        tmp.close()
        target = compose_service.slugify(target_id) if target_id else None
        try:
            return backup_service.restore_backup(
                tmp.name, target_id=target, overwrite=overwrite, restore_volumes=restore_volumes,
            )
        except backup_service.BackupError as exc:
            code = 409 if "already exists" in str(exc) else 400
            raise HTTPException(status_code=code, detail=str(exc)) from exc
    finally:
        if os.path.exists(tmp.name):
            os.unlink(tmp.name)


@app.get("/agent/health")
def health() -> dict:
    return {"status": "ok"}