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