homelab/lab/cli.py

import logging
import sys
import os

import pprint
from ansible.parsing.dataloader import DataLoader
from ansible.inventory.manager import InventoryManager
import subprocess

pp = pprint.PrettyPrinter()
logger = logging.getLogger(__name__)

DEFAULT_HOST_FILE = "hosts"
DEFAULT_PLAYBOOK = "playbook.yml"
DEFAULT_ANSIBLE_DIR = "lab/ansible"


def get_inventory() -> dict:
    host_file = os.path.join(DEFAULT_ANSIBLE_DIR, DEFAULT_HOST_FILE)
    loader = DataLoader()
    inventory = InventoryManager(loader=loader, sources=host_file)
    x = {}
    ignore = ("all", "ungrouped")
    x.update(inventory.groups["all"].serialize()["vars"])
    group_dict = inventory.get_groups_dict()

    for group in inventory.groups:
        if group in ignore:
            continue
        x.update({group: group_dict[group]})
    return x


def roles() -> list:
    sorted_roles = sorted(list(get_inventory().keys()))
    if sys.stdin and sys.stdin.isatty():
        for role in sorted_roles:
            print(role)
    return sorted_roles


def update(roles):
    if len(roles) == 1:
        roles = roles[0]
    else:
        roles = ",".join(roles)

    cmd = "ansible-playbook"
    limit = f"--limit={roles}"
    call_list = [cmd, DEFAULT_PLAYBOOK, "-i", DEFAULT_HOST_FILE, limit]
    print(call_list)
    if sys.stdin and sys.stdin.isatty():
        subprocess.run(call_list, cwd=DEFAULT_ANSIBLE_DIR)


def newjail(jail_name) -> None:
    if len(jail_name) == 1:
        jail_name = jail_name[0]
    else:
        print("Requires a single new jail name")
        return

    call_list = [
        "ansible",
        "rhea",
        "-m",
        "ansible.builtin.shell",
        "-a",
        f'"newjail {jail_name}"',
        "--become",
        "-K",
    ]
    print(call_list)
    if sys.stdin and sys.stdin.isatty():
        subprocess.run(call_list, cwd=DEFAULT_ANSIBLE_DIR)


def main():
    del sys.argv[0]

    if len(sys.argv) > 0:
        func = sys.argv.pop(0)
        args = sys.argv
        if args:
            return eval(func)(args)
        else:
            return eval(func)()


if __name__ == "__main__":
    main()