facts_to_aliases.py

from collections import defaultdict
from dataclasses import dataclass, field
from pprint import pprint
from typing import Any, Dict, Optional, Sequence


HOSTS = {
    "host123.example.org": {
        "enc": {
            "location": "zrh",
            "role": "web",
            "stage": "prod",
        },
    },
    "host141.example.org": {
        "enc": {
            "location": "zrh",
            "role": "web",
            "stage": "prod",
        },
    },
    "host176.example.org": {
        "enc": {
            "location": "zrh",
            "role": "web",
            "stage": "dev",
        },
    },
    "host223.example.org": {
        "enc": {
            "location": "zrh",
            "role": "mail",
        },
    },
    "host251.example.org": {
        "enc": {
            "location": "zrh",
            "role": "mail",
        },
    },
    "host298.example.org": {
        "enc": {
            "location": "zrh",
            "role": "dns",
        },
    },
    "host305.example.org": {
        "enc": {
            "location": "zrh",
            "role": "dns",
        },
    },
    "host390.example.org": {
        "enc": {
            "location": "win",
            "role": "dns",
        },
    },
    "host404.example.org": {
        "enc": {
            "location": "win",
            "role": "web",
            "stage": "dev",
        },
    },
    "host436.example.org": {
        "enc": {
            "location": "chr",
            "role": "web",
            "stage": "dev",
        },
    },
}


@dataclass
class HostsTreeNode:
    """Simple tree to aggregate hosts by hierarchy levels"""

    children: Optional[dict] = field(default_factory=dict)
    hosts: Optional[list] = field(default_factory=list)

    def host_labels(self, prefix: str = "enc") -> Dict[str, list]:
        """Hierarchy level based, indexed labels for node and child nodes grouped by host"""

        labels = defaultdict(list)
        for index, host in enumerate(self.hosts):
            labels[host].append(f"{prefix}.{index}")
        for level, child in self.children.items():
            for host, child_labels in child.host_labels(f"{prefix}.{level}").items():
                labels[host].extend(child_labels)
        return dict(labels)


def host_aliases(
    hosts: Dict[str, Dict[str, Any]], hierarchies: Sequence[Sequence[str]]
) -> Dict[str, list]:
    """Generate hierarchy based aliases for hosts"""

    root = HostsTreeNode()
    for hierarchy in hierarchies:
        for host, facts in hosts.items():
            level = root
            for enc_name in hierarchy:
                if (enc_value := facts["enc"].get(enc_name, None)) is None:
                    continue
                level = level.children.setdefault(enc_value, HostsTreeNode())
            level.hosts.append(host)
    return root.host_labels()


if __name__ == "__main__":
    pprint(
        host_aliases(
            hosts=HOSTS, hierarchies=(("location", "role", "stage"), ("role", "stage"))
        )
    )