FabledSteward/steward/ansible/executor.py

# steward/ansible/executor.py
from __future__ import annotations
import asyncio
import json
import logging
import os
import shutil
import tempfile
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from quart import Quart

logger = logging.getLogger(__name__)

_OUTPUT_CAP_BYTES = 1024 * 1024  # 1 MB DB cap
_FLUSH_LINES = 50
_FLUSH_SECS = 5.0


@dataclass
class _Done:
    status: str


# Per-run broadcast state (lives until process restarts)
_run_lines: dict[str, list[str]] = {}          # all output lines received so far
_run_done: dict[str, _Done] = {}               # set when run completes
_run_listeners: dict[str, list[asyncio.Queue]] = {}  # active SSE client queues


def register_listener(run_id: str) -> asyncio.Queue:
    """Create a listener queue for a run. Replays existing lines immediately."""
    q: asyncio.Queue = asyncio.Queue()
    for line in _run_lines.get(run_id, []):
        q.put_nowait(line)
    done = _run_done.get(run_id)
    if done:
        q.put_nowait(done)
    else:
        _run_listeners.setdefault(run_id, []).append(q)
    return q


def deregister_listener(run_id: str, q: asyncio.Queue) -> None:
    listeners = _run_listeners.get(run_id, [])
    if q in listeners:
        listeners.remove(q)


def _broadcast(run_id: str, item: str | _Done) -> None:
    if isinstance(item, str):
        _run_lines.setdefault(run_id, []).append(item)
    else:
        _run_done[run_id] = item
    for q in _run_listeners.get(run_id, []):
        q.put_nowait(item)
    if isinstance(item, _Done):
        _run_listeners.pop(run_id, None)


def build_ansible_command(
    playbook_path: str,
    inventory_path: str,
    params: dict | None = None,
) -> list[str]:
    """Build the ansible-playbook argv from a run's optional params.

    Pure (no IO). Everything is passed as argv — never shell-interpolated — so
    extra-var values and limits can contain arbitrary characters safely.
    params keys: extra_vars (list of "key=value"), limit, tags, check (bool).
    """
    cmd = ["ansible-playbook", playbook_path, "-i", inventory_path]
    params = params or {}
    for ev in params.get("extra_vars") or []:
        cmd += ["-e", ev]
    if params.get("limit"):
        cmd += ["--limit", params["limit"]]
    if params.get("tags"):
        cmd += ["--tags", params["tags"]]
    if params.get("check"):
        cmd += ["--check", "--diff"]
    return cmd


def build_credentials(creds: dict | None, tmpdir: str) -> tuple[list[str], list[tuple[str, str]]]:
    """Given global Ansible creds + a temp dir, return (extra argv, files to write).

    Pure: decides flags + file contents; the caller writes the files (mode 0600).
    Secrets are always passed via files — never on argv / the process list.
    creds keys: ssh_private_key, vault_password, become_password.
    """
    args: list[str] = []
    files: list[tuple[str, str]] = []
    creds = creds or {}

    key = creds.get("ssh_private_key")
    if key:
        path = os.path.join(tmpdir, "id_key")
        files.append((path, key if key.endswith("\n") else key + "\n"))
        args += ["--private-key", path]

    vault = creds.get("vault_password")
    if vault:
        path = os.path.join(tmpdir, "vault_pass")
        files.append((path, vault + "\n"))
        args += ["--vault-password-file", path]

    become = creds.get("become_password")
    if become:
        path = os.path.join(tmpdir, "become.yml")
        # A JSON string is valid YAML; keeps the password out of argv.
        files.append((path, "ansible_become_password: " + json.dumps(become) + "\n"))
        args += ["-e", "@" + path]

    return args, files


def ansible_env(creds: dict | None, base_env) -> dict:
    """Return a copy of base_env with ANSIBLE_HOST_KEY_CHECKING set from creds."""
    env = dict(base_env)
    env["ANSIBLE_HOST_KEY_CHECKING"] = "True" if (creds or {}).get("host_key_checking") else "False"
    return env


async def start_run(
    app: "Quart",
    run_id: str,
    playbook_path: str,
    inventory_path: str,
    source_path: str,
    params: dict | None = None,
) -> None:
    """Execute ansible-playbook as a subprocess and update the DB run row."""
    _run_lines[run_id] = []

    db_output = ""
    truncated = False
    lines_since_flush = 0
    last_flush_at = time.monotonic()

    async def _flush(final_status: str | None = None) -> None:
        nonlocal lines_since_flush, last_flush_at
        from sqlalchemy import update
        from steward.models.ansible import AnsibleRun, AnsibleRunStatus

        values: dict = {"output": db_output}
        if final_status is not None:
            values["status"] = AnsibleRunStatus(final_status)
            values["finished_at"] = datetime.now(timezone.utc)

        async with app.db_sessionmaker() as session:
            async with session.begin():
                await session.execute(
                    update(AnsibleRun).where(AnsibleRun.id == run_id).values(**values)
                )
        lines_since_flush = 0
        last_flush_at = time.monotonic()

    cmd = build_ansible_command(playbook_path, inventory_path, params)
    cwd = source_path if Path(source_path).exists() else None
    creds = app.config.get("ANSIBLE", {})

    # Materialize credentials into a private temp dir; removed in finally.
    tmpdir = tempfile.mkdtemp(prefix="steward-ansible-")
    os.chmod(tmpdir, 0o700)
    try:
        cred_args, cred_files = build_credentials(creds, tmpdir)
        for cred_path, content in cred_files:
            with open(cred_path, "w", encoding="utf-8") as cf:
                cf.write(content)
            os.chmod(cred_path, 0o600)

        proc = await asyncio.create_subprocess_exec(
            *cmd, *cred_args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
            cwd=cwd,
            env=ansible_env(creds, os.environ),
        )
        assert proc.stdout is not None

        async for raw_line in proc.stdout:
            line = raw_line.decode(errors="replace").rstrip("\n\r")
            _broadcast(run_id, line)

            if not truncated:
                candidate = (db_output + "\n" + line) if db_output else line
                if len(candidate.encode()) > _OUTPUT_CAP_BYTES:
                    db_output += "\n[output truncated]"
                    truncated = True
                else:
                    db_output = candidate

            lines_since_flush += 1
            elapsed = time.monotonic() - last_flush_at
            if lines_since_flush >= _FLUSH_LINES or elapsed >= _FLUSH_SECS:
                await _flush()

        await proc.wait()
        final_status = "success" if proc.returncode == 0 else "failed"

    except Exception:
        logger.exception("Ansible run %s failed with exception", run_id)
        final_status = "failed"
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)

    await _flush(final_status)
    _broadcast(run_id, _Done(status=final_status))