# 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))