kvmcli

#!/bin/python
import argparse
import logging
import subprocess
import multiprocessing as mp

from config import config
from src.info import report
from src.init import create_template
from src.images import copy_image
from src.args_gen import load_yaml_data, create_virt_install_args


logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")


def is_data_valid(data):
    return data is not None


def get_number_of_processes():
    return mp.cpu_count if config.FORK == 0 else config.FORK


def run_virt_install(virt_install_args):
    """Run virt-install command with arguments provided"""

    args = ["virt-install"]
    for key, value in virt_install_args.items():
        args.append(f"--{key}")
        if value:
            args.append(value)
    try:
        subprocess.run(args)
    except subprocess.CalledProcessError as e:
        logging.error(f"virt-install command failed: {e}")


def process_vm(vm, vm_index):
    """
    Process a single VM definition from the YAML data.
    """
    name = vm["info"].get("name", f"{config.IMAGE_NAME}-{vm_index+1}")

    if name in ignore:
        return

    image = copy_image(name, vm)

    if image["dest_image_exists"]:
        logging.info(f"Provisioning a new VM named {name}\n")
        vm_args = create_virt_install_args(vm_index, vm)
        run_virt_install(vm_args)


def apply(yaml_path):
    """
    Load a YAML file from the given file path, parse its content, and provision
    all the VM listed in it.
    """

    num_processes = get_number_of_processes()
    data = load_yaml_data(yaml_path)

    if not is_data_valid:
        logging.error("Invalid YAML data")
        exit()

    vm_index_tuples = [(vm, vm_index) for vm_index, vm in enumerate(data["vms"])]

    with mp.Pool(processes=num_processes) as pool:
        pool.starmap(process_vm, (vm_index_tuples))

    logging.info("All VMs provisioned successfully!")


def main():
    parser = argparse.ArgumentParser(
        prog="kvmcli",
        description="A Python script for managing virtual machines in a KVM-based environment.",
        epilog="Enjoy",
    )

    parser.add_argument(
        "-I", "--info", action="store_true", help="Print information about your cluster"
    )
    parser.add_argument(
        "-i", "--init", action="store_true", help="Create template file"
    )
    parser.add_argument(
        "-a", "--apply", action="store_true", help="apply configuration from YAML_FILE"
    )
    parser.add_argument("-f", "--file", metavar="YAML_FILE", help="Specify a yaml file")
    parser.add_argument("--ignore", metavar="NODE_NAME", help="Ignore NODE NAME")
    parser.add_argument("-v", "--version", action="store_true", help="Print version")

    args = parser.parse_args()

    if args.init:
        create_template()

    if args.file:
        yaml_file = args.file
    else:
        yaml_file = config.YAML_PATH

    if args.ignore:
        global ignore
        ignore = [node.strip() for node in args.ignore.split(",")]
    else:
        ignore = []

    if args.info:
        report(yaml_file)

    if args.apply:
        apply(yaml_file)

    if args.version:
        print(f"kvmcli {config.VERSION}")


if __name__ == "__main__":
    main()