FabledSteward/steward/ansible/inventory_gen.py

"""Ansible --list inventory generation from DB records."""
from __future__ import annotations

import re

_UNSAFE = re.compile(r"[^a-z0-9_]+")


def ansible_group_name(name: str) -> str:
    """A Steward group's name as Ansible accepts it (#5334): lower case, each
    run of anything but letters, digits and _ made one _, never starting with
    a digit. "Docker Swarm" is docker_swarm, which a play can name in
    `hosts:`; the display name stays as typed everywhere else."""
    safe = _UNSAFE.sub("_", name.lower()).strip("_") or "group"
    return f"g_{safe}" if safe[0].isdigit() else safe


def _group_names(targets: list) -> dict[str, str]:
    """{display name: Ansible name} for every group these targets are in. Two
    names that clean to the same thing ("Web servers", "web-servers") get _2,
    _3 in name order, so neither group's hosts or vars fold into the other's."""
    out: dict[str, str] = {}
    taken: set[str] = set()
    for name in sorted({g.name for t in targets for g in t.groups}):
        base = safe = ansible_group_name(name)
        n = 1
        while safe in taken or safe in ("all", "ungrouped", "_meta"):
            n += 1
            safe = f"{base}_{n}"
        taken.add(safe)
        out[name] = safe
    return out


def generate_inventory(targets: list) -> dict:
    """Build an Ansible --list JSON inventory from a list of AnsibleTarget objects.

    Accepts any objects with .name, .address, .ansible_vars (dict), .groups (list
    of objects with .name and .ansible_vars). Uses duck typing so unit tests can
    pass SimpleNamespace objects.

    Var precedence (lowest → highest, matching Ansible default):
      group vars (sorted alphabetically by group name) → target ansible_vars → ansible_host
    """
    inv: dict = {"all": {"hosts": []}, "_meta": {"hostvars": {}}}
    names = _group_names(targets)

    for target in targets:
        inv["all"]["hosts"].append(target.name)

        # Merge vars: group vars alphabetically, then host vars, then force ansible_host.
        merged: dict = {}
        for group in sorted(target.groups, key=lambda g: g.name):
            merged.update(group.ansible_vars or {})
        merged.update(target.ansible_vars or {})
        merged["ansible_host"] = target.address

        inv["_meta"]["hostvars"][target.name] = merged

        for group in target.groups:
            key = names[group.name]
            if key not in inv:
                inv[key] = {
                    "hosts": [],
                    "vars": dict(group.ansible_vars or {}),
                }
            inv[key]["hosts"].append(target.name)

    return inv


def inventory_to_yaml(inv: dict) -> str:
    """Convert the ``--list`` dict from generate_inventory into a STATIC YAML
    inventory string.

    The ``--list`` shape (``all.hosts`` is a *list*, vars under ``_meta``) is only
    valid as the stdout of an executable dynamic-inventory script. When written to
    a plain file and passed via ``-i``, Ansible's yaml plugin requires
    ``all.hosts`` to be a *dict* keyed by hostname — otherwise it fails to parse
    and silently falls back to implicit localhost. This emits that valid form.
    """
    import yaml
    hostvars = (inv.get("_meta") or {}).get("hostvars") or {}
    all_hosts = (inv.get("all") or {}).get("hosts") or []
    root: dict = {"hosts": {h: (hostvars.get(h) or {}) for h in all_hosts}}
    children: dict = {}
    for key, val in inv.items():
        if key in ("all", "_meta") or not isinstance(val, dict):
            continue
        grp: dict = {"hosts": {h: {} for h in (val.get("hosts") or [])}}
        if val.get("vars"):
            grp["vars"] = val["vars"]
        children[key] = grp
    if children:
        root["children"] = children
    return yaml.safe_dump({"all": root}, default_flow_style=False, sort_keys=True)


async def fetch_scope_targets(db, scope: str) -> list:
    """Query AnsibleTarget objects for a given inventory scope string.

    Scopes:
      steward:all              — every target
      steward:group:<uuid>     — targets in that group
      steward:target:<uuid>    — single target by ID

    Returns a list with groups pre-loaded (selectinload). Returns [] for
    repo:* scopes (caller handles those via legacy path).
    """
    from sqlalchemy import select
    from sqlalchemy.orm import selectinload
    from steward.models.ansible_inventory import AnsibleTarget, AnsibleGroup

    if scope == "steward:all":
        stmt = (
            select(AnsibleTarget)
            .options(selectinload(AnsibleTarget.groups))
            .order_by(AnsibleTarget.name)
        )
    elif scope.startswith("steward:group:"):
        group_id = scope[len("steward:group:"):]
        stmt = (
            select(AnsibleTarget)
            .join(AnsibleTarget.groups)
            .where(AnsibleGroup.id == group_id)
            .options(selectinload(AnsibleTarget.groups))
            .order_by(AnsibleTarget.name)
        )
    elif scope.startswith("steward:target:"):
        target_id = scope[len("steward:target:"):]
        stmt = (
            select(AnsibleTarget)
            .where(AnsibleTarget.id == target_id)
            .options(selectinload(AnsibleTarget.groups))
        )
    else:
        return []

    result = await db.execute(stmt)
    return list(result.scalars().all())


def describe_scope(scope: str, targets: dict[str, str], groups: dict[str, str]) -> str:
    """A run's inventory scope in words: "target web-1", "group Docker Swarm",
    "all inventory targets", "<source>/<inventory file>". `targets` and
    `groups` map id → name; a scope naming a deleted one says so."""
    scope = scope or ""
    if scope == "steward:all":
        return "all inventory targets"
    for prefix, kind, names in (("steward:target:", "target", targets),
                                ("steward:group:", "group", groups)):
        if scope.startswith(prefix):
            ident = scope[len(prefix):]
            name = names.get(ident)
            return f"{kind} {name}" if name else f"{kind} {ident[:8]} (deleted)"
    if scope.startswith("repo:"):
        source, _, inv = scope[len("repo:"):].rpartition(":")
        return f"{source}/{inv}" if source else inv
    return scope


def run_title(playbook: str, scope_label: str | None) -> str:
    """A run named by what it ran where: "update.yml on web-1",
    "site.yml on group Docker Swarm", "ping.yml on all inventory targets"."""
    where = scope_label or ""
    if where.startswith("target "):
        where = where[len("target "):]
    return f"{playbook} on {where}" if where else playbook


async def scope_labels(db, scopes) -> dict[str, str]:
    """{scope: describe_scope(scope)} for the given scopes, in two queries."""
    from sqlalchemy import select
    from steward.models.ansible_inventory import AnsibleGroup, AnsibleTarget

    scopes = {s for s in scopes if s}
    t_ids = {s[len("steward:target:"):] for s in scopes if s.startswith("steward:target:")}
    g_ids = {s[len("steward:group:"):] for s in scopes if s.startswith("steward:group:")}
    targets = dict((await db.execute(
        select(AnsibleTarget.id, AnsibleTarget.name).where(AnsibleTarget.id.in_(t_ids))
    )).all()) if t_ids else {}
    groups = dict((await db.execute(
        select(AnsibleGroup.id, AnsibleGroup.name).where(AnsibleGroup.id.in_(g_ids))
    )).all()) if g_ids else {}
    return {s: describe_scope(s, targets, groups) for s in scopes}