1#!/usr/bin/env python
2
3"""
4The minionswarm script will start a group of salt minions with different ids
5on a single system to test scale capabilities
6"""
7# pylint: disable=resource-leakage
8
9import hashlib
10import optparse
11import os
12import random
13import shutil
14import signal
15import subprocess
16import sys
17import tempfile
18import time
19import uuid
20
21import salt
22import salt.utils.files
23import salt.utils.yaml
24import tests.support.runtests
25
26OSES = [
27    "Arch",
28    "Ubuntu",
29    "Debian",
30    "CentOS",
31    "Fedora",
32    "Gentoo",
33    "AIX",
34    "Solaris",
35]
36VERS = [
37    "2014.1.6",
38    "2014.7.4",
39    "2015.5.5",
40    "2015.8.0",
41]
42
43
44def parse():
45    """
46    Parse the cli options
47    """
48    parser = optparse.OptionParser()
49    parser.add_option(
50        "-m",
51        "--minions",
52        dest="minions",
53        default=5,
54        type="int",
55        help="The number of minions to make",
56    )
57    parser.add_option(
58        "-M",
59        action="store_true",
60        dest="master_too",
61        default=False,
62        help="Run a local master and tell the minions to connect to it",
63    )
64    parser.add_option(
65        "--master",
66        dest="master",
67        default="salt",
68        help="The location of the salt master that this swarm will serve",
69    )
70    parser.add_option(
71        "--name",
72        "-n",
73        dest="name",
74        default="ms",
75        help=(
76            "Give the minions an alternative id prefix, this is used "
77            "when minions from many systems are being aggregated onto "
78            "a single master"
79        ),
80    )
81    parser.add_option(
82        "--rand-os",
83        dest="rand_os",
84        default=False,
85        action="store_true",
86        help="Each Minion claims a different os grain",
87    )
88    parser.add_option(
89        "--rand-ver",
90        dest="rand_ver",
91        default=False,
92        action="store_true",
93        help="Each Minion claims a different version grain",
94    )
95    parser.add_option(
96        "--rand-machine-id",
97        dest="rand_machine_id",
98        default=False,
99        action="store_true",
100        help="Each Minion claims a different machine id grain",
101    )
102    parser.add_option(
103        "--rand-uuid",
104        dest="rand_uuid",
105        default=False,
106        action="store_true",
107        help="Each Minion claims a different UUID grain",
108    )
109    parser.add_option(
110        "-k",
111        "--keep-modules",
112        dest="keep",
113        default="",
114        help="A comma delimited list of modules to enable",
115    )
116    parser.add_option(
117        "-f",
118        "--foreground",
119        dest="foreground",
120        default=False,
121        action="store_true",
122        help="Run the minions with debug output of the swarm going to the terminal",
123    )
124    parser.add_option(
125        "--temp-dir",
126        dest="temp_dir",
127        default=None,
128        help="Place temporary files/directories here",
129    )
130    parser.add_option(
131        "--no-clean",
132        action="store_true",
133        default=False,
134        help="Don't cleanup temporary files/directories",
135    )
136    parser.add_option(
137        "--root-dir",
138        dest="root_dir",
139        default=None,
140        help="Override the minion root_dir config",
141    )
142    parser.add_option(
143        "--transport",
144        dest="transport",
145        default="zeromq",
146        help="Declare which transport to use, default is zeromq",
147    )
148    parser.add_option(
149        "--start-delay",
150        dest="start_delay",
151        default=0.0,
152        type="float",
153        help="Seconds to wait between minion starts",
154    )
155    parser.add_option(
156        "-c",
157        "--config-dir",
158        default="",
159        help="Pass in a configuration directory containing base configuration.",
160    )
161    parser.add_option("-u", "--user", default=tests.support.runtests.this_user())
162
163    options, _args = parser.parse_args()
164
165    opts = {}
166
167    for key, val in options.__dict__.items():
168        opts[key] = val
169
170    return opts
171
172
173class Swarm:
174    """
175    Create a swarm of minions
176    """
177
178    def __init__(self, opts):
179        self.opts = opts
180
181        # If given a temp_dir, use it for temporary files
182        if opts["temp_dir"]:
183            self.swarm_root = opts["temp_dir"]
184        else:
185            # If given a root_dir, keep the tmp files there as well
186            if opts["root_dir"]:
187                tmpdir = os.path.join(opts["root_dir"], "tmp")
188            else:
189                tmpdir = opts["root_dir"]
190            self.swarm_root = tempfile.mkdtemp(
191                prefix="mswarm-root", suffix=".d", dir=tmpdir
192            )
193
194        if self.opts["transport"] == "zeromq":
195            self.pki = self._pki_dir()
196        self.zfill = len(str(self.opts["minions"]))
197
198        self.confs = set()
199
200        random.seed(0)
201
202    def _pki_dir(self):
203        """
204        Create the shared pki directory
205        """
206        path = os.path.join(self.swarm_root, "pki")
207        if not os.path.exists(path):
208            os.makedirs(path)
209
210            print("Creating shared pki keys for the swarm on: {}".format(path))
211            subprocess.call(
212                "salt-key -c {0} --gen-keys minion --gen-keys-dir {0} "
213                "--log-file {1} --user {2}".format(
214                    path,
215                    os.path.join(path, "keys.log"),
216                    self.opts["user"],
217                ),
218                shell=True,
219            )
220            print("Keys generated")
221        return path
222
223    def start(self):
224        """
225        Start the magic!!
226        """
227        if self.opts["master_too"]:
228            master_swarm = MasterSwarm(self.opts)
229            master_swarm.start()
230        minions = MinionSwarm(self.opts)
231        minions.start_minions()
232        print("Starting minions...")
233        # self.start_minions()
234        print("All {} minions have started.".format(self.opts["minions"]))
235        print("Waiting for CTRL-C to properly shutdown minions...")
236        while True:
237            try:
238                time.sleep(5)
239            except KeyboardInterrupt:
240                print("\nShutting down minions")
241                self.clean_configs()
242                break
243
244    def shutdown(self):
245        """
246        Tear it all down
247        """
248        print("Killing any remaining running minions")
249        subprocess.call('pkill -KILL -f "python.*salt-minion"', shell=True)
250        if self.opts["master_too"]:
251            print("Killing any remaining masters")
252            subprocess.call('pkill -KILL -f "python.*salt-master"', shell=True)
253        if not self.opts["no_clean"]:
254            print("Remove ALL related temp files/directories")
255            shutil.rmtree(self.swarm_root)
256        print("Done")
257
258    def clean_configs(self):
259        """
260        Clean up the config files
261        """
262        for path in self.confs:
263            pidfile = "{}.pid".format(path)
264            try:
265                try:
266                    with salt.utils.files.fopen(pidfile) as fp_:
267                        pid = int(fp_.read().strip())
268                    os.kill(pid, signal.SIGTERM)
269                except ValueError:
270                    pass
271                if os.path.exists(pidfile):
272                    os.remove(pidfile)
273                if not self.opts["no_clean"]:
274                    shutil.rmtree(path)
275            except OSError:
276                pass
277
278
279class MinionSwarm(Swarm):
280    """
281    Create minions
282    """
283
284    def start_minions(self):
285        """
286        Iterate over the config files and start up the minions
287        """
288        self.prep_configs()
289        for path in self.confs:
290            cmd = "salt-minion -c {} --pid-file {}".format(path, "{}.pid".format(path))
291            if self.opts["foreground"]:
292                cmd += " -l debug &"
293            else:
294                cmd += " -d &"
295            subprocess.call(cmd, shell=True)
296            time.sleep(self.opts["start_delay"])
297
298    def mkconf(self, idx):
299        """
300        Create a config file for a single minion
301        """
302        data = {}
303        if self.opts["config_dir"]:
304            spath = os.path.join(self.opts["config_dir"], "minion")
305            with salt.utils.files.fopen(spath) as conf:
306                data = salt.utils.yaml.safe_load(conf) or {}
307        minion_id = "{}-{}".format(self.opts["name"], str(idx).zfill(self.zfill))
308
309        dpath = os.path.join(self.swarm_root, minion_id)
310        if not os.path.exists(dpath):
311            os.makedirs(dpath)
312
313        data.update(
314            {
315                "id": minion_id,
316                "user": self.opts["user"],
317                "cachedir": os.path.join(dpath, "cache"),
318                "master": self.opts["master"],
319                "log_file": os.path.join(dpath, "minion.log"),
320                "grains": {},
321            }
322        )
323
324        if self.opts["transport"] == "zeromq":
325            minion_pkidir = os.path.join(dpath, "pki")
326            if not os.path.exists(minion_pkidir):
327                os.makedirs(minion_pkidir)
328                minion_pem = os.path.join(self.pki, "minion.pem")
329                minion_pub = os.path.join(self.pki, "minion.pub")
330                shutil.copy(minion_pem, minion_pkidir)
331                shutil.copy(minion_pub, minion_pkidir)
332            data["pki_dir"] = minion_pkidir
333        elif self.opts["transport"] == "tcp":
334            data["transport"] = "tcp"
335
336        if self.opts["root_dir"]:
337            data["root_dir"] = self.opts["root_dir"]
338
339        path = os.path.join(dpath, "minion")
340
341        if self.opts["keep"]:
342            keep = self.opts["keep"].split(",")
343            modpath = os.path.join(os.path.dirname(salt.__file__), "modules")
344            fn_prefixes = (fn_.partition(".")[0] for fn_ in os.listdir(modpath))
345            ignore = [fn_prefix for fn_prefix in fn_prefixes if fn_prefix not in keep]
346            data["disable_modules"] = ignore
347
348        if self.opts["rand_os"]:
349            data["grains"]["os"] = random.choice(OSES)
350        if self.opts["rand_ver"]:
351            data["grains"]["saltversion"] = random.choice(VERS)
352        if self.opts["rand_machine_id"]:
353            data["grains"]["machine_id"] = hashlib.md5(minion_id).hexdigest()
354        if self.opts["rand_uuid"]:
355            data["grains"]["uuid"] = str(uuid.uuid4())
356
357        with salt.utils.files.fopen(path, "w+") as fp_:
358            salt.utils.yaml.safe_dump(data, fp_)
359        self.confs.add(dpath)
360
361    def prep_configs(self):
362        """
363        Prepare the confs set
364        """
365        for idx in range(self.opts["minions"]):
366            self.mkconf(idx)
367
368
369class MasterSwarm(Swarm):
370    """
371    Create one or more masters
372    """
373
374    def __init__(self, opts):
375        super().__init__(opts)
376        self.conf = os.path.join(self.swarm_root, "master")
377
378    def start(self):
379        """
380        Prep the master start and fire it off
381        """
382        # sys.stdout for no newline
383        sys.stdout.write("Generating master config...")
384        self.mkconf()
385        print("done")
386
387        sys.stdout.write("Starting master...")
388        self.start_master()
389        print("done")
390
391    def start_master(self):
392        """
393        Do the master start
394        """
395        cmd = "salt-master -c {} --pid-file {}".format(
396            self.conf, "{}.pid".format(self.conf)
397        )
398        if self.opts["foreground"]:
399            cmd += " -l debug &"
400        else:
401            cmd += " -d &"
402        subprocess.call(cmd, shell=True)
403
404    def mkconf(self):  # pylint: disable=W0221
405        """
406        Make a master config and write it'
407        """
408        data = {}
409        if self.opts["config_dir"]:
410            spath = os.path.join(self.opts["config_dir"], "master")
411            with salt.utils.files.fopen(spath) as conf:
412                data = salt.utils.yaml.safe_load(conf)
413        data.update(
414            {
415                "log_file": os.path.join(self.conf, "master.log"),
416                "open_mode": True,  # TODO Pre-seed keys
417            }
418        )
419
420        os.makedirs(self.conf)
421        path = os.path.join(self.conf, "master")
422
423        with salt.utils.files.fopen(path, "w+") as fp_:
424            salt.utils.yaml.safe_dump(data, fp_)
425
426    def shutdown(self):
427        print("Killing master")
428        subprocess.call('pkill -KILL -f "python.*salt-master"', shell=True)
429        print("Master killed")
430
431
432# pylint: disable=C0103
433if __name__ == "__main__":
434    swarm = Swarm(parse())
435    try:
436        swarm.start()
437    finally:
438        swarm.shutdown()
439