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
import shutil
from dataclasses import asdict

import tempfile

import json

from fastapi import (
    Depends,
    FastAPI,
    File,
    Form,
    Header,
    HTTPException,
    Query,
    Request,
    UploadFile,
    WebSocket,
    WebSocketDisconnect,
)
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,
    device_service,
    file_service,
    image_service,
    network_service,
    stats_service,
    update_service,
    volume_service,
)

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

# Map network_service's DockerError codes to HTTP status. forbidden is mapped to
# 400 (not 403) so the central proxy doesn't misread it as a token failure.
_DOCKER_STATUS = {"invalid_request": 400, "forbidden": 400, "not_found": 404}


def _map_docker(exc: DockerError):
    code = _DOCKER_STATUS.get(exc.error)
    if code:
        raise HTTPException(status_code=code, detail=exc.detail or exc.error)
    raise exc  # falls through to the global 502 DockerError handler

AGENT_VERSION = "0.23.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


class NetworkCreateBody(BaseModel):
    name: str
    driver: str = "bridge"
    subnet: str | None = None
    gateway: str | None = None
    internal: bool = False
    attachable: bool = True


class ContainerRefBody(BaseModel):
    container: str
    aliases: list[str] | None = None
    force: bool = False


class FileWriteBody(BaseModel):
    path: str
    content: str


class FileNameBody(BaseModel):
    path: str
    name: str


class FileRenameBody(BaseModel):
    path: str
    new_name: str


class FileTransferBody(BaseModel):
    src: str
    dest_dir: str
    overwrite: bool = False


def _file_guard(fn, *args, **kwargs):
    try:
        return fn(*args, **kwargs)
    except file_service.BrowseError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


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


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


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


def _mem_info() -> tuple[int, int]:
    """Return (total_bytes, used_bytes) from meminfo (used = total - available)."""
    for base in (settings.HOST_PROC_PATH, "/proc"):
        try:
            vals: dict[str, int] = {}
            with open(os.path.join(base, "meminfo"), "r", encoding="utf-8") as fh:
                for line in fh:
                    parts = line.split(":")
                    if len(parts) == 2 and parts[0] in ("MemTotal", "MemAvailable", "MemFree"):
                        try:
                            vals[parts[0]] = int(parts[1].split()[0]) * 1024  # kB -> bytes
                        except ValueError:
                            pass
            total = vals.get("MemTotal", 0)
            available = vals.get("MemAvailable", vals.get("MemFree", 0))
            return total, max(total - available, 0)
        except OSError:
            continue
    return 0, 0


def _disk_info() -> tuple[int, int]:
    """Return (total_bytes, used_bytes) for the host disk backing the stacks dir."""
    for path in (settings.STACKS_DIR, "/"):
        try:
            usage = shutil.disk_usage(path)
            return usage.total, usage.used
        except OSError:
            continue
    return 0, 0


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})"
    mem_total, mem_used = _mem_info()
    disk_total, disk_used = _disk_info()
    return {
        "hostname": _hostname(),
        "docker_version": docker_version,
        "host_os": host_os,
        "cpu_cores": os.cpu_count() or 0,
        "mem_total": mem_total,
        "mem_used": mem_used,
        "disk_total": disk_total,
        "disk_used": disk_used,
        "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]:
    try:
        summaries = compose_service.stack_status_summaries()
    except DockerError:
        summaries = {}
    return [_summary(sid, summaries) for sid in compose_service.discover_stacks()]


@app.get("/agent/stacks/stats", dependencies=[Depends(verify_token)])
def stacks_stats() -> dict:
    return stats_service.stack_stats()


@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:
        raw = compose_service.containers_for_stack(stack_id)
        containers = [asdict(c) for c in raw]
        status = compose_service.compute_status(stack_id, raw)
    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)


# --------------------------------------------------------------------------- #
# Networks
# --------------------------------------------------------------------------- #


@app.get("/agent/networks", dependencies=[Depends(verify_token)])
def list_networks() -> list[dict]:
    return network_service.list_networks()


@app.get("/agent/networks/{network_id}", dependencies=[Depends(verify_token)])
def inspect_network(network_id: str) -> dict:
    try:
        return network_service.inspect_network(network_id)
    except DockerError as exc:
        _map_docker(exc)


@app.get("/agent/networks/{network_id}/containers", dependencies=[Depends(verify_token)])
def network_containers(network_id: str) -> list[dict]:
    try:
        return network_service.connectable_containers(network_id)
    except DockerError as exc:
        _map_docker(exc)


@app.post("/agent/networks/{network_id}/connect", dependencies=[Depends(verify_token)])
def connect_container(network_id: str, body: ContainerRefBody) -> dict:
    try:
        network_service.connect_container(network_id, body.container, body.aliases)
    except DockerError as exc:
        _map_docker(exc)
    return {"ok": True}


@app.post("/agent/networks/{network_id}/disconnect", dependencies=[Depends(verify_token)])
def disconnect_container(network_id: str, body: ContainerRefBody) -> dict:
    try:
        network_service.disconnect_container(network_id, body.container, body.force)
    except DockerError as exc:
        _map_docker(exc)
    return {"ok": True}


@app.post("/agent/networks", dependencies=[Depends(verify_token)], status_code=201)
def create_network(body: NetworkCreateBody) -> dict:
    try:
        return network_service.create_network(body.model_dump())
    except DockerError as exc:
        _map_docker(exc)


@app.delete("/agent/networks/{network_id}", dependencies=[Depends(verify_token)])
def delete_network(network_id: str) -> dict:
    try:
        network_service.delete_network(network_id)
    except DockerError as exc:
        _map_docker(exc)
    return {"ok": True}


@app.post("/agent/networks/prune", dependencies=[Depends(verify_token)])
def prune_networks() -> dict:
    return network_service.prune_networks()


# --------------------------------------------------------------------------- #
# Images
# --------------------------------------------------------------------------- #


@app.get("/agent/images", dependencies=[Depends(verify_token)])
def list_images() -> list[dict]:
    return image_service.list_images()


@app.get("/agent/images/updates", dependencies=[Depends(verify_token)])
def image_updates() -> dict:
    return update_service.get_cache()


@app.post("/agent/images/check", dependencies=[Depends(verify_token)])
async def image_check() -> dict:
    return await update_service.check_all()


# --------------------------------------------------------------------------- #
# Volumes
# --------------------------------------------------------------------------- #


@app.get("/agent/volumes", dependencies=[Depends(verify_token)])
def list_volumes() -> list[dict]:
    return volume_service.list_volumes()


@app.get("/agent/volumes/sizes", dependencies=[Depends(verify_token)])
def volume_sizes(force: bool = Query(False)) -> dict:
    return volume_service.volume_sizes(force=force)


@app.delete("/agent/volumes/{name}", dependencies=[Depends(verify_token)])
def delete_volume(name: str, force: bool = Query(False)) -> dict:
    vols = {v["name"]: v for v in volume_service.list_volumes()}
    if name in vols and vols[name]["in_use"] and not force:
        raise HTTPException(
            status_code=409,
            detail={
                "error": "volume_in_use",
                "detail": f"Volume '{name}' is used by: {', '.join(vols[name]['used_by'])}",
            },
        )
    volume_service.remove_volume(name, force=force)
    return {"ok": True}


@app.post("/agent/volumes/prune", dependencies=[Depends(verify_token)])
def prune_volumes() -> dict:
    return volume_service.prune_volumes()


# --------------------------------------------------------------------------- #
# File browser (sandboxed by this agent's ALLOWED_BROWSE_ROOTS/HOST_ROOT_PREFIX)
# --------------------------------------------------------------------------- #


@app.get("/agent/files/list", dependencies=[Depends(verify_token)])
def files_list(path: str = Query("/"), show_hidden: bool = Query(False)) -> dict:
    return _file_guard(device_service.browse, path, show_hidden)


@app.get("/agent/files/read", dependencies=[Depends(verify_token)])
def files_read(path: str = Query(...)) -> dict:
    return _file_guard(file_service.read_file, path)


@app.get("/agent/files/download", dependencies=[Depends(verify_token)])
def files_download(path: str = Query(...)):
    real, filename = _file_guard(file_service.resolve_download, path)
    return FileResponse(real, filename=filename, media_type="application/octet-stream")


@app.put("/agent/files/write", dependencies=[Depends(verify_token)])
def files_write(body: FileWriteBody) -> dict:
    return _file_guard(file_service.write_file, body.path, body.content)


@app.post("/agent/files/mkdir", dependencies=[Depends(verify_token)])
def files_mkdir(body: FileNameBody) -> dict:
    return _file_guard(file_service.create_dir, body.path, body.name)


@app.post("/agent/files/touch", dependencies=[Depends(verify_token)])
def files_touch(body: FileNameBody) -> dict:
    return _file_guard(file_service.create_file, body.path, body.name)


@app.post("/agent/files/rename", dependencies=[Depends(verify_token)])
def files_rename(body: FileRenameBody) -> dict:
    return _file_guard(file_service.rename, body.path, body.new_name)


@app.post("/agent/files/copy", dependencies=[Depends(verify_token)])
def files_copy(body: FileTransferBody) -> dict:
    return _file_guard(file_service.copy, body.src, body.dest_dir, body.overwrite)


@app.post("/agent/files/move", dependencies=[Depends(verify_token)])
def files_move(body: FileTransferBody) -> dict:
    return _file_guard(file_service.move, body.src, body.dest_dir, body.overwrite)


@app.delete("/agent/files", dependencies=[Depends(verify_token)])
def files_delete(path: str = Query(...), recursive: bool = Query(False)) -> dict:
    return _file_guard(file_service.delete, path, recursive)


@app.post("/agent/files/upload", dependencies=[Depends(verify_token)])
async def files_upload(
    path: str = Form(...),
    overwrite: bool = Form(False),
    rel_path: str = Form(""),
    file: UploadFile = File(...),
) -> dict:
    real = _file_guard(
        file_service.upload_target, path, file.filename or "", overwrite, rel_path or None
    )
    tmp = tempfile.NamedTemporaryFile(delete=False, dir=os.path.dirname(real))
    try:
        while chunk := await file.read(1024 * 1024):
            tmp.write(chunk)
        tmp.close()
        os.replace(tmp.name, real)
    except OSError as exc:
        if os.path.exists(tmp.name):
            os.unlink(tmp.name)
        raise HTTPException(status_code=400, detail=f"Upload failed: {exc}") from exc
    return {"ok": True, "name": rel_path or file.filename}


@app.websocket("/agent/ws/logs/{stack_id}")
async def ws_logs(websocket: WebSocket, stack_id: str, token: str | None = Query(default=None)):
    """Stream `docker compose logs -f` to the central app (token via query param)."""
    await websocket.accept()
    expected = settings.AGENT_TOKEN
    if not expected or token != expected:
        await websocket.close(code=4401)
        return
    if not os.path.isdir(compose_service.stack_dir(stack_id)):
        await websocket.send_text(json.dumps({"type": "error", "detail": "stack not found"}))
        await websocket.close()
        return
    args = ["logs", "--no-color", "--tail", "200", "--timestamps", "-f"]
    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": None, "line": line})
            )
    except WebSocketDisconnect:
        pass
    except Exception as exc:  # noqa: BLE001
        try:
            await websocket.send_text(json.dumps({"type": "error", "detail": str(exc)}))
        except Exception:  # noqa: BLE001
            pass


@app.websocket("/agent/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 to the central app so the
    browser sees deploy progress live (token via query param)."""
    await websocket.accept()
    expected = settings.AGENT_TOKEN
    if not expected or token != expected:
        await websocket.close(code=4401)
        return
    if not os.path.isdir(compose_service.stack_dir(stack_id)):
        await websocket.send_text(json.dumps({"type": "error", "detail": "stack not found"}))
        await websocket.close()
        return
    compose_service.mark_busy(stack_id)
    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:
                await websocket.send_text(json.dumps({"type": "done", "returncode": payload}))
    except WebSocketDisconnect:
        # Browser navigated away; the compose subprocess keeps running.
        pass
    except Exception as exc:  # noqa: BLE001
        try:
            await websocket.send_text(json.dumps({"type": "error", "detail": str(exc)}))
        except Exception:  # noqa: BLE001
            pass
    finally:
        compose_service.clear_busy(stack_id)


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