#!/usr/bin/python

##Autore: Nicolo' Palazzini
##Liberamente ispirato a labcalcoloctl di Elisa Aliverti (all rights reserved)
##Ultima modifica: Maggio 2019

from time import time
start = time()
import argparse
import os
import subprocess, sys
from threading import Thread

choices = ('status','start','stop')
parser = argparse.ArgumentParser(description='Simple tool to handle xfreerdp on LCM nodes.', usage=sys.argv[0]+' {'+','.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.1', 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()

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.status()

    # 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.up:
            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 status(self):
        statuscmd = "ps aux | grep 'xfreerdp' | grep -v grep"
        ssh = self.sshcommand(statuscmd) 
        result = [ l for l in ssh.stdout.readlines() if '/usr/bin/xfreerdp /cert-ignore /f /v:tolab.fisica.unimi.it' in l ]

        if result == []:
            return False
        else:
            return True

    def xstart(self):
        if self.up:
            if not self.running:
                startcmd = "startx /usr/bin/xfreerdp /cert-ignore /f /v:tolab.fisica.unimi.it &> /dev/null &"
                self.sshcommand(startcmd)
                print 'xfreerdp started on ' + self.hostname
            else:
                print 'xfreerdp is already running on ' + self.hostname
        else:
            print self.hostname + ' is not up.'

    def xstop(self):
        if self.up:
            if self.running:
                # Try to kill sending SIGTERM, wait, check if running and if so issue SIGKILL.
                # (Do this asynchronously, on the remote node)
                stopcmd = '''PCMDL="/usr/bin/xfreerdp /cert-ignore /f /v:tolab.fisica.unimi.it"
                             pkill -xf "$PCMDL"
                             sleep 5
                             pgrep -xf "$PCMDL" && \
                             pkill -9 -xf "$PCMDL"'''
                self.sshcommand(stopcmd)
                print 'xfreerdp stopped on ' + self.hostname
            else:
                print 'xfreerdp is not running on ' + self.hostname
        else:
            print self.hostname + ' is not up.'

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 xfreerdp on LCM nodes."
	print '\nTry: "'+sys.argv[0]+' --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 ]

for i in nodes :
    i.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 "xfreerdp is running on:"
        for i in running : print '\t',i
    else:
        print "xfreerdp is not running"
    if len(down):
        print "Down nodes:"
        for i in down :  print '\t',i

elif args.cmd == 'start':
    for i in nodes: i.xstart()

elif args.cmd == 'stop':
    for i in nodes: i.xstop()