evil-minions

#!/usr/bin/python3
'''
Script to start a salt-minion that will act as many minions (from salt-master's
point of view).
Every response sent by the minion will be replicated and sent out separately
to the master with different minion ids, as if multiple minions were connected
and responding in the same way.
'''

import argparse
import distutils.spawn
import logging
from multiprocessing import Process, cpu_count, Semaphore
import sys

import salt.log

from evilminions.proxy import start_proxy
from evilminions.vampire import Vampire
from evilminions.hydra import Hydra

def main():
    # parse commandline switches
    parser = argparse.ArgumentParser(description='Starts a salt-minion and many evil minions that mimic it.')
    parser.add_argument('--count', dest='count', type=int, default=10,
                       help='number of evil minions (default: 10)')
    parser.add_argument('--id-prefix', dest='prefix', default='evil',
                       help='minion id prefix for evil minions. (default: evil)')
    parser.add_argument('--id-offset', dest='offset', type=int, default=0,
                       help='minion id counter offset for evil minions. (default: 0)')
    parser.add_argument('--ramp-up-delay', dest='ramp_up_delay', type=int, default=0,
                       help='time between evil minion starts in seconds (default: 0)')
    parser.add_argument('--slowdown-factor', dest='slowdown_factor', type=float, default=0.0,
                       help='slow down evil minions (default is 0.0 "as fast as possible", ' +
                       '1.0 is "as fast as the original")')
    parser.add_argument('--random-slowdown-factor', dest='random_slowdown_factor', type=float, default=0.0,
                       help='a random extra delay expressed as a percentage of the original time (default: 0.0)')
    parser.add_argument('--processes', dest='processes', type=int, default=cpu_count(),
                       help='number of concurrent processes (default is the CPU count: %d)' % cpu_count())
    parser.add_argument('--keysize', dest='keysize', type=int, default=2048,
                       help='size of Salt keys generated for each evil minion (default: 2048)')
    args, remaining_args = parser.parse_known_args()

    # set up logging
    salt.log.setup_console_logger(log_level='debug')
    log = logging.getLogger(__name__)
    log.debug("Starting evil-minions, setting up infrastructure...")

    # set up Hydras, one per process. Hydras are so called because they have
    # many HydraHeads, each HydraHead running one evil-minion
    semaphore = Semaphore(0)
    chunks = minion_chunks(args.count, args.processes)
    hydras = [Process(target=Hydra(i).start, kwargs={
        'hydra_count': len(chunks),
        'chunk': chunk,
        'prefix': args.prefix,
        'offset': args.offset,
        'ramp_up_delay': args.ramp_up_delay,
        'slowdown_factor': args.slowdown_factor,
        'random_slowdown_factor': args.random_slowdown_factor,
        'keysize': args.keysize,
        'semaphore': semaphore
    }) for i, chunk in enumerate(chunks)]

    for hydra in hydras:
        hydra.start()

    for hydra in hydras:
        semaphore.acquire()

    # set up Vampire, to intercept messages from the original minion and send them to Proxy
    vampire = Vampire()
    vampire.attach()

    # set up Proxy, to route messages from Vampire to Hydras
    proxy = Process(target=start_proxy, kwargs={'semaphore': semaphore})
    proxy.start()
    semaphore.acquire()

    # restore salt-minion's orignal command line args and launch it
    sys.argv = [sys.argv[0]] + remaining_args
    salt_minion_executable = distutils.spawn.find_executable('salt-minion')
    exec(compile(open(salt_minion_executable).read(), salt_minion_executable, 'exec'))

def minion_chunks(count, processes):
    '''Returns chunks of minion indexes, per process (eg. 4 minions and 3 processes: [[0, 1], [2], [3]])'''
    minions_per_process = count // processes
    rest = count % processes
    lengths = [minions_per_process + 1] * rest + [minions_per_process] * (processes - rest)
    starts = [sum(lengths[:i]) for i in range(processes)]

    indexes = list(range(count))

    return [indexes[starts[i]:starts[i] + lengths[i]] for i in range(processes)]

if __name__ == "__main__":
    main()