#!/usr/bin/env python3

##Author: Elisa Aliverti
##Last edit: January 2020 - Matteo Bonamassa

from time import time
start = 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(description="Simple tool to handle LabCalcolo's VMs.", usage='%(prog)s {'+','.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 2.0',
                     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
		self.running  = False
		self.up       = False

	# Run method called on Thread start. Check if host is up and if is running a VM
	def run(self):
		if self.isup() :
			self.up = True
			self.running=self.vmstatus()

	# 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 vmstatus(self):
		statuscmd = "ps aux | grep qemu | grep -v grep"
		ssh = self.sshcommand(statuscmd)
		result = [ l for l in ssh.stdout.readlines() if 'qemu' in l ]
		if result == []:
			return False
		else:
			return True

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

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


	def vmstop(self):
		if self.up:
			if self.running:
				stopcmd = main_path + " 0"
				self.sshcommand(stopcmd)
				print('VM stopped on ' + self.hostname)
			else:
				print('VM is not running on ' + self.hostname)
		else:
			print(self.hostname + ' is not up.')

	def vmdoctor(self) :
		to_search=['qemu', 'spicec']
		# It's ugly (and more or less useless), 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.up:
			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)

			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 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 = []

# Show usage if no arguments
if len(sys.argv) < 2:
	parser.print_usage()
	print("\nSimple tool to handle LabCalcolo's VMs.")
	print('\nTry: "labcalcoloctl --help" to display help message.')
	sys.exit(1)

# 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:
    nodes = [ j for j in Hosts if j.hostname in args.node ]

# Start the threads and run commands on nodes
for n in nodes :
    n.start()

running=[]
down=[]

num=len(nodes)
index=0
print(' Querying ' + str(num) + ' hosts...')
for i in nodes:
	# Rejoin them when their work is done
	i.join()
	if i.running:
		running.append(i.hostname)
	if not i.up :
		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 args.cmd == 'status':
    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)

#########################################################################################################
## Since lcm2 nodes have no VMs installed a control is added to make sure operator knows what he's doing.
## When future (or present?) admins will fix the issue this part should be deleted.

elif args.cmd == 'start':
	yes = {'yes','y','ye'}
	no = {'no','n'}
	control = True
	parsed = False
	for n in nodes:
		if n.location == 'LCM2' and parsed == False :
			print('You are trying to start a VM on lcm2 nodes. Are you sure you want to continue? [yes/no]')
			while True:
				choice = input().lower()
				if choice in yes:
					parsed = True
					break
				elif choice in no:
					control = False
					parsed = True
					break
				else:
					sys.stdout.write("Be kind, respond with 'yes' or 'no'")

	if control == True:
		for i in nodes: i.vmstart()

#########################################################################################################
## Uncomment this part when the lcm2 VMs issue will be fixed.

#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()