Python for Automation: using pdsh for a menu-driven command execution environment

I’ve been playing around with some quick system automation scripts that are handy to use when you don’t want / need to setup a chef or puppet action. I like to keep all of my hostnames and login details in a MySQL database (a cmdb actually) but for this example we’ll just use a couple of nested lists. This script executes commands in parallel across the hosts you choose in the menu system via the “pdsh” command, so make sure you have that installed before running. Alternately you can change the command call to use ssh instead of pdsh for a serialized execution, but that’s not as fun or fast. With some customizations here and there you can expand this to operate parallelized jobs for simplifying daily work in database administration, usage reporting, log file parsing, or other system automation as you see fit. Here’s the code. Comments welcome as always!

#!/usr/bin/env python
## NAME: menu_parallel_execution.py
## DATE: 2011-02-04
## AUTHOR: Matt Reid
## WEBSITE: http://kontrollsoft.com
## EMAIL: mreid@kontrollsoft.com
## LICENSE: BSD http://www.opensource.org/licenses/bsd-license.php

import commands
import sys
import os
import operator

ver = sys.version.split(' ')[0].split(".")
major=ver[:1]
minor=ver[1:2]
version="%s.%s"%(major[0],minor[0])
if version in ('2.4','2.3','2.2','2.1','2.0'):
	pyver = "old"
else:
	from subprocess import Popen, PIPE, STDOUT, call
	pyver = "new"

def main():
        #This section of host definitions could be replaced by a database
        #call if you are storing your hosts in that manner, but we'll keep
        #this script simple for the time being.
	dc0 = [
		["host0-0","10.0.0.1","ssh-username"],
		["host0-1","10.0.0.2","ss-username"]]

	dc1 = [
		["host1-0","10.1.0.1","ssh-username"],
		["host1-1","10.1.0.2","ss-username"]]

	dc2 = [
		["host2-0","10.2.0.1","ssh-username"],
		["host2-1","10.2.0.2","ss-username"]]

	print '''[1] Datacenter-1
[2] Datacenter-2
[3] Datacenter-3
'''
	dc = int(raw_input("Datacenter ID: "))
	if dc == 1:
		hosts = dc0
	elif dc == 2:
		hosts = dc1
	elif dc == 3:
		hosts = dc3

	nodes = []
	stay = True
	while stay == True:
		i = 0
		nest = sorted(hosts, key=operator.itemgetter(0))
		for x in range(len(nest)):
			print "[%i] %s | %s"%(i,nest[i][1],nest[i][0])
			i+=1

		print "\nSelect node to add to execution list"
		ch = int(raw_input("ID: "))
		xx = [nest[ch][1],nest[ch][2]]
		nodes.append(xx)
		s = str(raw_input("\nAdd another node? [Y/n] "))
		if s == "n" or s == "N":
			stay = False

	if(pyver == "new"):
		addrs = ""
		for node in nodes:
			address = node[1]+"@"+node[0]+","
			address = address.rstrip("\n")
			addrs = addrs+address

		addrs = addrs.strip(",")
		cmd = str(raw_input("\nEnter the command to execute: "))				

		try:
			c = "pdsh -w %s %s"%(addrs,cmd)
			print "Executing: %s"%(c)
			call(c,shell=True)
		except:
			print "Failed to execute pdsh command: %s"%(c)
			sys.exit(1)

		sys.exit(0)

	if(pyver == "old"):
		print "Please upgrade to Python 2.6+"
		sys.exit(1)

## START
if __name__ == "__main__":
	try:
		retval = main()
	except (KeyboardInterrupt, SystemExit):
		sys.exit(1)