Display ansible inventories with Python and the rich console library

inv.py

"""
Display Ansible inventories as a tree of environments, groups and hosts as leaves.

A nice diversion/excuse to dive into the rich console library.
It will expand simple host ranges, allowing grep commands to find
group names for given hosts.

Usage:
  ./inv.py
  ./inv.py | grep -B10 some-hostname
  
Example:
                                                                                                                                                                                                                   
    ...                                                                                                                                                                                                            
    └──  experimental                                                                                                                                                                                              
        ├──  🌏 tomcat9                                                                                                                                                                                            
        │   └── 📂 tomcat9_experimental                                                                                                                                                                            
        │       └──  💻  experimental-tomcat9-1.domain.com                                                                                                                                                         
        ├──  🌏 couchbase                                                                                                                                                                                          
        │   ├── 📂 couchbase_experimental                                                                                                                                                                          
        │   │   ├──  💻  1.2.3.4         
        │   │   └──  💻  5.6.7.8              
        │   └── 📂 index_experimental               
        │       └── 📂 experimental_index_1                          
        │           └──  💻  experimental-index-1.domain.com
        ├──  🌏 elasticsearch              
        │   └── 📂 es_experimental                           
        │       ├── 💻  experimental-es-1.domain.com         
        │       └── 💻  experimental-es-2.domain.com
        └──  🌏 services                         
            └── 📂 services_experimental                
                ├──  💻  e-services1.domain.com         
                └──  💻  e-services2.domain.com

"""

from pathlib import Path
import re
import yaml

from rich import print as rprint
from rich.tree import Tree


def read_inventory(filename) -> dict:
    with open(filename, 'r') as inventory_file:
        groups = yaml.load(inventory_file)
        return groups


def build_inventory_tree(envs: dict) -> Tree:
    tree = Tree("Ansible Inventories")

    for env in envs.keys():
        env_node = tree.add(f"[bold cyan] {env}")

        for group in envs[env].keys():
            group_node = env_node.add(f"[bold blue] 🌏 {group}")
            walk_inventory(envs[env][group], group_node)

    return tree


def walk_inventory(node, tree: Tree) -> None:
    if type(node) != dict:
        return

    is_group = "children" in node
    is_leaf = "hosts" in node

    if is_group:
        children = node["children"]
        for c in children.keys():
            subgroup = tree.add(f"📂 {c}")
            walk_inventory(children[c], subgroup)
    elif is_leaf:
        hosts_spec = node["hosts"]

        if type(hosts_spec) not in (list, iter):
            hosts_spec = [hosts_spec]

        expanded_hosts = expand_hosts_range(hosts_spec)

        for host in expanded_hosts:
            tree.add(f"[bold white] 💻  {host}")


def expand_hosts_range(hosts_spec):
    expanded_hosts = []
    range_re = re.compile(r'^([a-z0-9\._-]+)\[([0-9]+):([0-9]+)\]([a-z0-9\._-]+)$')

    for host in hosts_spec:
        if range_re.match(host):
            groups = range_re.match(host)
            prefix = groups[1]
            suffix = groups[4]
            m = int(groups[2])
            n = int(groups[3])
            for i in range(m, m + n):
                expanded_hosts.append(prefix + str(i) + suffix)
        else:
            expanded_hosts.append(host)

    return expanded_hosts


p = Path('./inventories')

envs = dict()
inventories = list(p.glob('./**/inventory.yml'))

for i in inventories:

    env = i.parent.stem
    groups = read_inventory(i)

    envs[env] = groups


rprint(build_inventory_tree(envs))