FabledSteward/steward/ansible/routes.py

# steward/ansible/routes.py
from __future__ import annotations
import asyncio
import uuid
from datetime import datetime, timezone

from quart import (
    Blueprint, Response, current_app, render_template,
    request, session,
)
from sqlalchemy import select

from steward.ansible import executor, sources as src_module
from steward.auth.middleware import require_role
from steward.models.ansible import AnsibleRun, AnsibleRunStatus
from steward.models.users import User, UserRole

ansible_bp = Blueprint("ansible", __name__, url_prefix="/ansible")


def _get_sources() -> list[dict]:
    ansible_cfg = current_app.config.get("ANSIBLE", {})
    return src_module.get_sources(ansible_cfg)


@ansible_bp.get("/")
@require_role(UserRole.viewer)
async def index():
    async with current_app.db_sessionmaker() as db:
        result = await db.execute(
            select(AnsibleRun).order_by(AnsibleRun.started_at.desc()).limit(50)
        )
        runs = result.scalars().all()
    return await render_template("ansible/index.html", runs=runs)


@ansible_bp.get("/browse")
@require_role(UserRole.viewer)
async def browse():
    all_sources = _get_sources()
    source_data = []
    for source in all_sources:
        playbooks = src_module.discover_playbooks(source["path"])
        inventories = src_module.discover_inventories(source["path"])
        source_data.append({
            "source": source,
            "playbooks": playbooks,
            "inventories": inventories,
        })
    return await render_template("ansible/browse.html", source_data=source_data)


@ansible_bp.get("/browse/<source_name>/<path:playbook_path>")
@require_role(UserRole.viewer)
async def view_playbook(source_name: str, playbook_path: str):
    all_sources = _get_sources()
    source = next((s for s in all_sources if s["name"] == source_name), None)
    if source is None:
        return "Source not found", 404
    contents = src_module.read_playbook(source["path"], playbook_path)
    if contents is None:
        return "Playbook not found", 404
    return await render_template(
        "ansible/browse.html",
        source_data=[],
        view_source=source_name,
        view_path=playbook_path,
        view_contents=contents,
    )


@ansible_bp.post("/runs")
@require_role(UserRole.operator)
async def create_run():
    form = await request.form
    playbook_path = form.get("playbook_path", "").strip()
    source_name = form.get("source_name", "").strip()
    inventory_path = form.get("inventory_path", "").strip()

    if not playbook_path or not inventory_path:
        return "playbook_path and inventory_path are required", 400

    all_sources = _get_sources()
    source = next((s for s in all_sources if s["name"] == source_name), None)
    if source is None:
        return "Source not found", 404

    # Optional run parameters (all passed as argv by the executor — no shell).
    extra_vars: list[str] = []
    for line in (form.get("extra_vars", "") or "").splitlines():
        line = line.strip()
        if not line:
            continue
        if "=" not in line:
            return f"Invalid extra var (expected key=value): {line!r}", 400
        extra_vars.append(line)
    limit = (form.get("limit", "") or "").strip()
    tags = (form.get("tags", "") or "").strip()
    check = "check" in form

    params: dict = {}
    if extra_vars:
        params["extra_vars"] = extra_vars
    if limit:
        params["limit"] = limit
    if tags:
        params["tags"] = tags
    if check:
        params["check"] = True
    params_or_none = params or None

    run_id = str(uuid.uuid4())
    now = datetime.now(timezone.utc)

    run = AnsibleRun(
        id=run_id,
        playbook_path=playbook_path,
        inventory_path=inventory_path,
        source_name=source_name,
        triggered_by=session["user_id"],
        status=AnsibleRunStatus.running,
        started_at=now,
        params=params_or_none,
    )
    async with current_app.db_sessionmaker() as db:
        async with db.begin():
            db.add(run)

    task = asyncio.create_task(
        executor.start_run(
            current_app._get_current_object(),  # type: ignore[attr-defined]
            run_id,
            playbook_path,
            inventory_path,
            source["path"],
            params_or_none,
        )
    )
    task.add_done_callback(
        lambda t: t.exception() and current_app.logger.error(
            "Ansible run %s raised: %s", run_id, t.exception()
        )
    )

    return await render_template(
        "ansible/run_started.html",
        run=run,
        source=source,
    )


@ansible_bp.get("/runs/<run_id>/stream")
@require_role(UserRole.viewer)
async def stream_run(run_id: str):
    async with current_app.db_sessionmaker() as db:
        result = await db.execute(select(AnsibleRun).where(AnsibleRun.id == run_id))
        run = result.scalar_one_or_none()
    if run is None:
        return "Not found", 404

    async def generate():
        if run.status != AnsibleRunStatus.running:
            yield f"event: done\ndata: {run.status.value}\n\n"
            return

        q = executor.register_listener(run_id)
        try:
            while True:
                msg = await q.get()
                if isinstance(msg, executor._Done):
                    yield f"event: done\ndata: {msg.status}\n\n"
                    break
                # Escape newlines in SSE data field
                safe = msg.replace("\n", " ")
                yield f"event: output\ndata: {safe}\n\n"
        finally:
            executor.deregister_listener(run_id, q)

    return Response(
        generate(),
        content_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )


@ansible_bp.get("/runs/<run_id>")
@require_role(UserRole.viewer)
async def run_detail(run_id: str):
    async with current_app.db_sessionmaker() as db:
        result = await db.execute(select(AnsibleRun).where(AnsibleRun.id == run_id))
        run = result.scalar_one_or_none()
        if run is None:
            return "Not found", 404
        # NULL triggered_by = an automated (system) run; otherwise resolve the user.
        triggered_label = "system"
        if run.triggered_by:
            res = await db.execute(
                select(User.username).where(User.id == run.triggered_by))
            triggered_label = res.scalar_one_or_none() or "(deleted user)"
    return await render_template(
        "ansible/run_detail.html", run=run, triggered_label=triggered_label)