salt/tests/minionswarm.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
The minionswarm script will start a group of salt minions with different ids
on a single system to test scale capabilities
"""
# pylint: disable=resource-leakage
# Import Python Libs
from __future__ import absolute_import, print_function

import hashlib
import optparse
import os
import random
import shutil
import signal
import subprocess
import sys
import tempfile
import time
import uuid

# Import salt libs
import salt
import salt.utils.files
import salt.utils.yaml
import tests.support.runtests

# Import third party libs
from salt.ext import six
from salt.ext.six.moves import range  # pylint: disable=import-error,redefined-builtin

OSES = [
    "Arch",
    "Ubuntu",
    "Debian",
    "CentOS",
    "Fedora",
    "Gentoo",
    "AIX",
    "Solaris",
]
VERS = [
    "2014.1.6",
    "2014.7.4",
    "2015.5.5",
    "2015.8.0",
]


def parse():
    """
    Parse the cli options
    """
    parser = optparse.OptionParser()
    parser.add_option(
        "-m",
        "--minions",
        dest="minions",
        default=5,
        type="int",
        help="The number of minions to make",
    )
    parser.add_option(
        "-M",
        action="store_true",
        dest="master_too",
        default=False,
        help="Run a local master and tell the minions to connect to it",
    )
    parser.add_option(
        "--master",
        dest="master",
        default="salt",
        help="The location of the salt master that this swarm will serve",
    )
    parser.add_option(
        "--name",
        "-n",
        dest="name",
        default="ms",
        help=(
            "Give the minions an alternative id prefix, this is used "
            "when minions from many systems are being aggregated onto "
            "a single master"
        ),
    )
    parser.add_option(
        "--rand-os",
        dest="rand_os",
        default=False,
        action="store_true",
        help="Each Minion claims a different os grain",
    )
    parser.add_option(
        "--rand-ver",
        dest="rand_ver",
        default=False,
        action="store_true",
        help="Each Minion claims a different version grain",
    )
    parser.add_option(
        "--rand-machine-id",
        dest="rand_machine_id",
        default=False,
        action="store_true",
        help="Each Minion claims a different machine id grain",
    )
    parser.add_option(
        "--rand-uuid",
        dest="rand_uuid",
        default=False,
        action="store_true",
        help="Each Minion claims a different UUID grain",
    )
    parser.add_option(
        "-k",
        "--keep-modules",
        dest="keep",
        default="",
        help="A comma delimited list of modules to enable",
    )
    parser.add_option(
        "-f",
        "--foreground",
        dest="foreground",
        default=False,
        action="store_true",
        help=(
            "Run the minions with debug output of the swarm going to " "the terminal"
        ),
    )
    parser.add_option(
        "--temp-dir",
        dest="temp_dir",
        default=None,
        help="Place temporary files/directories here",
    )
    parser.add_option(
        "--no-clean",
        action="store_true",
        default=False,
        help="Don't cleanup temporary files/directories",
    )
    parser.add_option(
        "--root-dir",
        dest="root_dir",
        default=None,
        help="Override the minion root_dir config",
    )
    parser.add_option(
        "--transport",
        dest="transport",
        default="zeromq",
        help="Declare which transport to use, default is zeromq",
    )
    parser.add_option(
        "--start-delay",
        dest="start_delay",
        default=0.0,
        type="float",
        help="Seconds to wait between minion starts",
    )
    parser.add_option(
        "-c",
        "--config-dir",
        default="",
        help=("Pass in a configuration directory containing base configuration."),
    )
    parser.add_option("-u", "--user", default=tests.support.runtests.this_user())

    options, _args = parser.parse_args()

    opts = {}

    for key, val in six.iteritems(options.__dict__):
        opts[key] = val

    return opts


class Swarm(object):
    """
    Create a swarm of minions
    """

    def __init__(self, opts):
        self.opts = opts

        # If given a temp_dir, use it for temporary files
        if opts["temp_dir"]:
            self.swarm_root = opts["temp_dir"]
        else:
            # If given a root_dir, keep the tmp files there as well
            if opts["root_dir"]:
                tmpdir = os.path.join(opts["root_dir"], "tmp")
            else:
                tmpdir = opts["root_dir"]
            self.swarm_root = tempfile.mkdtemp(
                prefix="mswarm-root", suffix=".d", dir=tmpdir
            )

        if self.opts["transport"] == "zeromq":
            self.pki = self._pki_dir()
        self.zfill = len(str(self.opts["minions"]))

        self.confs = set()

        random.seed(0)

    def _pki_dir(self):
        """
        Create the shared pki directory
        """
        path = os.path.join(self.swarm_root, "pki")
        if not os.path.exists(path):
            os.makedirs(path)

            print("Creating shared pki keys for the swarm on: {0}".format(path))
            subprocess.call(
                "salt-key -c {0} --gen-keys minion --gen-keys-dir {0} "
                "--log-file {1} --user {2}".format(
                    path, os.path.join(path, "keys.log"), self.opts["user"],
                ),
                shell=True,
            )
            print("Keys generated")
        return path

    def start(self):
        """
        Start the magic!!
        """
        if self.opts["master_too"]:
            master_swarm = MasterSwarm(self.opts)
            master_swarm.start()
        minions = MinionSwarm(self.opts)
        minions.start_minions()
        print("Starting minions...")
        # self.start_minions()
        print("All {0} minions have started.".format(self.opts["minions"]))
        print("Waiting for CTRL-C to properly shutdown minions...")
        while True:
            try:
                time.sleep(5)
            except KeyboardInterrupt:
                print("\nShutting down minions")
                self.clean_configs()
                break

    def shutdown(self):
        """
        Tear it all down
        """
        print("Killing any remaining running minions")
        subprocess.call('pkill -KILL -f "python.*salt-minion"', shell=True)
        if self.opts["master_too"]:
            print("Killing any remaining masters")
            subprocess.call('pkill -KILL -f "python.*salt-master"', shell=True)
        if not self.opts["no_clean"]:
            print("Remove ALL related temp files/directories")
            shutil.rmtree(self.swarm_root)
        print("Done")

    def clean_configs(self):
        """
        Clean up the config files
        """
        for path in self.confs:
            pidfile = "{0}.pid".format(path)
            try:
                try:
                    with salt.utils.files.fopen(pidfile) as fp_:
                        pid = int(fp_.read().strip())
                    os.kill(pid, signal.SIGTERM)
                except ValueError:
                    pass
                if os.path.exists(pidfile):
                    os.remove(pidfile)
                if not self.opts["no_clean"]:
                    shutil.rmtree(path)
            except (OSError, IOError):
                pass


class MinionSwarm(Swarm):
    """
    Create minions
    """

    def start_minions(self):
        """
        Iterate over the config files and start up the minions
        """
        self.prep_configs()
        for path in self.confs:
            cmd = "salt-minion -c {0} --pid-file {1}".format(
                path, "{0}.pid".format(path)
            )
            if self.opts["foreground"]:
                cmd += " -l debug &"
            else:
                cmd += " -d &"
            subprocess.call(cmd, shell=True)
            time.sleep(self.opts["start_delay"])

    def mkconf(self, idx):
        """
        Create a config file for a single minion
        """
        data = {}
        if self.opts["config_dir"]:
            spath = os.path.join(self.opts["config_dir"], "minion")
            with salt.utils.files.fopen(spath) as conf:
                data = salt.utils.yaml.safe_load(conf) or {}
        minion_id = "{0}-{1}".format(self.opts["name"], str(idx).zfill(self.zfill))

        dpath = os.path.join(self.swarm_root, minion_id)
        if not os.path.exists(dpath):
            os.makedirs(dpath)

        data.update(
            {
                "id": minion_id,
                "user": self.opts["user"],
                "cachedir": os.path.join(dpath, "cache"),
                "master": self.opts["master"],
                "log_file": os.path.join(dpath, "minion.log"),
                "grains": {},
            }
        )

        if self.opts["transport"] == "zeromq":
            minion_pkidir = os.path.join(dpath, "pki")
            if not os.path.exists(minion_pkidir):
                os.makedirs(minion_pkidir)
                minion_pem = os.path.join(self.pki, "minion.pem")
                minion_pub = os.path.join(self.pki, "minion.pub")
                shutil.copy(minion_pem, minion_pkidir)
                shutil.copy(minion_pub, minion_pkidir)
            data["pki_dir"] = minion_pkidir
        elif self.opts["transport"] == "tcp":
            data["transport"] = "tcp"

        if self.opts["root_dir"]:
            data["root_dir"] = self.opts["root_dir"]

        path = os.path.join(dpath, "minion")

        if self.opts["keep"]:
            keep = self.opts["keep"].split(",")
            modpath = os.path.join(os.path.dirname(salt.__file__), "modules")
            fn_prefixes = (fn_.partition(".")[0] for fn_ in os.listdir(modpath))
            ignore = [fn_prefix for fn_prefix in fn_prefixes if fn_prefix not in keep]
            data["disable_modules"] = ignore

        if self.opts["rand_os"]:
            data["grains"]["os"] = random.choice(OSES)
        if self.opts["rand_ver"]:
            data["grains"]["saltversion"] = random.choice(VERS)
        if self.opts["rand_machine_id"]:
            data["grains"]["machine_id"] = hashlib.md5(minion_id).hexdigest()
        if self.opts["rand_uuid"]:
            data["grains"]["uuid"] = str(uuid.uuid4())

        with salt.utils.files.fopen(path, "w+") as fp_:
            salt.utils.yaml.safe_dump(data, fp_)
        self.confs.add(dpath)

    def prep_configs(self):
        """
        Prepare the confs set
        """
        for idx in range(self.opts["minions"]):
            self.mkconf(idx)


class MasterSwarm(Swarm):
    """
    Create one or more masters
    """

    def __init__(self, opts):
        super(MasterSwarm, self).__init__(opts)
        self.conf = os.path.join(self.swarm_root, "master")

    def start(self):
        """
        Prep the master start and fire it off
        """
        # sys.stdout for no newline
        sys.stdout.write("Generating master config...")
        self.mkconf()
        print("done")

        sys.stdout.write("Starting master...")
        self.start_master()
        print("done")

    def start_master(self):
        """
        Do the master start
        """
        cmd = "salt-master -c {0} --pid-file {1}".format(
            self.conf, "{0}.pid".format(self.conf)
        )
        if self.opts["foreground"]:
            cmd += " -l debug &"
        else:
            cmd += " -d &"
        subprocess.call(cmd, shell=True)

    def mkconf(self):  # pylint: disable=W0221
        """
        Make a master config and write it'
        """
        data = {}
        if self.opts["config_dir"]:
            spath = os.path.join(self.opts["config_dir"], "master")
            with salt.utils.files.fopen(spath) as conf:
                data = salt.utils.yaml.safe_load(conf)
        data.update(
            {
                "log_file": os.path.join(self.conf, "master.log"),
                "open_mode": True,  # TODO Pre-seed keys
            }
        )

        os.makedirs(self.conf)
        path = os.path.join(self.conf, "master")

        with salt.utils.files.fopen(path, "w+") as fp_:
            salt.utils.yaml.safe_dump(data, fp_)

    def shutdown(self):
        print("Killing master")
        subprocess.call('pkill -KILL -f "python.*salt-master"', shell=True)
        print("Master killed")


# pylint: disable=C0103
if __name__ == "__main__":
    swarm = Swarm(parse())
    try:
        swarm.start()
    finally:
        swarm.shutdown()