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 os
import time
import signal
import optparse
import subprocess
import random
import tempfile
import shutil
import sys
import hashlib
import uuid

# Import salt libs
import salt
import salt.utils.files
import salt.utils.yaml

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


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.helpers.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()