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'"/usr/local/bin/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()