labcalcoloctl/labcalcoloctl

#!/usr/bin/python

##Author: Elisa Aliverti
##Last edit: 22/02/2017 - Silva

from time import time
import argparse
import textwrap
import os
import subprocess, sys
from threading import Thread
from getpass import getpass

choices = ('status','start','stop', 'doctor') 
parser = argparse.ArgumentParser(usage='labcalcolo {'+','.join(choices)+'} [options]')

## positional arguments
parser.add_argument( 'cmd', nargs="?", choices=choices, default='status', 
                     help='Specify command to execute, default is \'status\'' )

## optional arguments
parser.add_argument( '-a', '--all', action='store_true', dest='lcm',
                     help='All LCM nodes are considered' )
parser.add_argument( '-n', nargs='+', dest='node',  
                     help='Select one or more nodes (at least one)' )
parser.add_argument( '-1', '--lcm1', action='store_true', dest='lcm1',
                     help='LCM1 nodes are considered' )
parser.add_argument( '-2', '--lcm2', action='store_true', dest='lcm2',
                     help='LCM2 nodes are considered' )
parser.add_argument( '-v', '--version', action='version', version='%(prog)s 1.3', 
                     help='Print program version' )	   
## 
args = parser.parse_args()

def print_progressbar(index, num) :
    sys.stdout.write('\r ['
                        + '='*index
                        + '>'*(1-int(index/num))
                        + ' '*(num-index-1) + ']')
    sys.stdout.flush()
 

## main path (you don't say)
main_path = "/var/etc/vmctl"

class Host(Thread):
   # Constructor
   def __init__(self, name, location):
      # Fork the thread first thing
      Thread.__init__(self)
      # Variables initialization
      self.hostname = name
      self.location = location

   # Ping the host to see if it's up
   def isup(self):
      # Is the host up?
      ping = os.popen("ping -w1 -c1 " + self.hostname, "r")
      # print("pinging "  self.hostname)
      if "0 received" in ping.read():
      	 return False
      else:
         return True

   def sshcommand(self, command):
      if self.isup():
         ssh = subprocess.Popen( ["ssh", "%s" % self.hostname, command],
       	         		 shell = False, 
                                 stdout = subprocess.PIPE,
                                 stderr = subprocess.PIPE )
         return ssh
      else:
         print self.hostname + ' is not up.'
         return False

   def vmstart(self):
      if not self.vmstatus():
          startcmd = main_path + " 1"
          self.sshcommand(startcmd)
          print 'VM is now starting on ' + self.hostname
      else:
          print 'VM is already running on ' + self.hostname
          
   def vmstop(self):
      stopcmd = main_path + " 0"
      self.sshcommand(stopcmd)

   def vmstatus(self):
      statuscmd = "ps aux | grep qemu | grep -v grep"
      if self.isup():
         ssh = self.sshcommand(statuscmd)
         result = [ l for l in ssh.stdout.readlines() if 'qemu' in l ]
         if result == []:
              return False
         else:
              return True
      else:
          return -1
    
   def vmdoctor(self) :
      to_search=['qemu', 'spicec']
      # It's ugly, but it should not be necessary: passwords should not be visible from ps aux
      pw_remove=['s/,password=\w*$//', 's/-w \w*$//']
      # Build the status query with programs names and relative pw remove strings
      status_query="ps aux | grep -E '" + '|'.join(to_search) + "' | grep -v grep | sed '" + ';'.join(pw_remove)+ "'"
      if self.isup():
        ssh = self.sshcommand(status_query).stdout.readlines()
        qemu_status = [ l for l in ssh if 'qemu' in l ] # Filter for a single command
        if len(qemu_status) :
            print "Qemu command running on", self.hostname+":\n", qemu_status 
        else :
            print "Qemu is not running on", self.hostname
            # Do we want to ask this? There's the start command to do this...
            #answer=raw_input("Do you want to start a VM now? [y/n] ")
            #if answer=="y" : self.vmstart()

        spicec_status = [ l for l in ssh if 'spicec' in l ]
        if not len(spicec_status) and len(qemu_status) :
            print "Spicec is not running on", self.hostname, 
            if raw_input("Do you want to start it now? [y/n] ")=="y" :
                pw=getpass("Vm password: ")
                spiceccmd="export DISPLAY=:4 ; spicec -f -h 127.0.0.1 -p 5900 -w "+pw+" &"
                self.sshcommand(spiceccmd) # How to check if everything went as expected?
                # Check if everything works fine now
                spicec_status = [ l for l in self.sshcommand(status_query).stdout.readlines() if 'spicec' in l ]
        if len(spicec_status) :
            print "Spicec command running on", self.hostname+": ", spicec_status

### end class Host

## Host list
Hosts = [
  Host('abe', 'LCM1'),
  Host('crash', 'LCM1'),
  Host('duke', 'LCM1'),
  Host('glados', 'LCM1'),
  Host('lara', 'LCM1'),
  Host('link', 'LCM1'),
  Host('king', 'LCM1'),
  Host('pang', 'LCM1'),
  Host('pong', 'LCM1'), 
  Host('snake', 'LCM1'),
  Host('sonic', 'LCM1'),
  Host('spyro', 'LCM1'),
  Host('yoshi', 'LCM1'),
  Host('actarus', 'LCM2'),
  Host('elwood', 'LCM2'),
  Host('gex', 'LCM2'),
  Host('gin', 'LCM2'),
  Host('jake', 'LCM2'),
  Host('kirk', 'LCM2'),
  Host('martini', 'LCM2'),
  Host('picard', 'LCM2'),
  Host('q', 'LCM2'),
  Host('raziel', 'LCM2'),
  Host('sarek',  'LCM2'),
  Host('spock',  'LCM2'),
  Host('tron',  'LCM2'),
  Host('worf',  'LCM2'),
  Host('zombie', 'LCM2')
]

nodes = []

# Filter hostlist according to arguments
if args.lcm:
    nodes  = Hosts
elif args.lcm1:
    nodes = [ host for host in Hosts if host.location == 'LCM1' ]
elif args.lcm2:
    nodes = [ host for host in Hosts if host.location == 'LCM2' ]
elif args.node:
    for i in args.node:
        for j in Hosts: 
             if (i==j.hostname): nodes.append(j)


# Run commands on nodes
running=[]
down=[]
if args.cmd == 'status':
    start = time()
    num=len(nodes)
    index=0
    print ' Querying ' + str(num) + ' hosts...'
    for i in nodes:
        if i.vmstatus():
          running.append(i.hostname)
        if i.vmstatus()<0 :
          down.append(i.hostname)
        index += 1
        print_progressbar(index, num)
    # New line after progress bar     
    print '\n Done... (%(t).3f s)' % {'t': (time() - start)}

    if len(running):
        print "VM(s) running on:"
        for i in running : print '\t',i
    else :
        print "No VMs are running"
    if len(down):
        print "Down nodes:"
        for i in down :  print '\t',i

elif args.cmd == 'start':
   for i in nodes: i.vmstart()
elif args.cmd == 'stop':
   for i in nodes: i.vmstop()
elif args.cmd == 'doctor' :
   for i in nodes : i.vmdoctor()