#!/usr/bin/python
#
# Python whoall: parallel, multi-threaded version
# First author: jacopogh - nov 2006
# Modified and updated by admins & 150 over time
#
# Requires python version 2.7 or above
# Please read the wiki before editing this file

from time import time
start = time()

# Import needed modules
import argparse
import collections
import os
import telnetlib
import textwrap
from threading import Thread
from socket import error
from sys import argv,exit,stdout

# Some array definitions
lcm1 = ['abe','crash','duke','glados','lara','link','king','pang','pong','snake','sonic','spyro','yoshi']
lcm2 = ['actarus','elwood','gex','gin','jake','kirk','martini','picard','q','raziel','sarek','spock','tron','worf','zombie']
laur = ['eskimo','orion','tilde']

math4 = ['sonic','crash','duke','raziel','actarus','gin','kirk','martini','picard','sarek','tron','worf','zombie','q','elwood','jake']
math5 = ['abe','glados','link','king','pang','pong','snake','spyro','yoshi']
math = math4 + math5

cuda = ['jacobi','yukawa','tesla']


# Important variables: gods, hosts and term colors
gods = ['root','andreatsh','andreamalerba','lorenzouboldi','stefanobalzan','eugeniothieme']
# Former admins: bother them at your own risk
chucknorris = [ 'agalli', 'ikki', 'buddino', 'alex', 'ema', 'ktf', 'davideg', 'jacopogh', 'lampo', 'gian', 'rbondesan', 'scolari', 'emanueleb', 'giani_matteo','gabryv','fran','alqahirah','giorgio_ruffa','palazzi','algebrato','blanc','blue','silviacotroneo'] 


parser = argparse.ArgumentParser( epilog='Please, report bugs or unwanted behavior to: working@lcm.mi.infn.it' )

## Optional arguments
parser.add_argument( '-1', '--lcm1', action='store_true', default=False, dest='lcm1', help='LCM1 hosts ' )
parser.add_argument( '-2', '--lcm2', action='store_true', default=False, dest='lcm2', help='LCM2 hosts ' )
parser.add_argument( '-l', '--laur', action='store_true', default=False, dest='laur', help='only LAUR hosts ' )
parser.add_argument( '-c', '--condor', action='store_true', default=False, dest='condor', help='condor nodes '  )
parser.add_argument( '-C', '--cuda', action='store_true', default=False, dest='cuda', help='only CUDA hosts ' )
parser.add_argument( '-f', '--full', action='store_true', default=False, dest='full', help='empty, unreachable and unavailable info' )
parser.add_argument( '-m', '--math', action='store_true', default=False, dest='math', help='only Mathematica hosts ' )
parser.add_argument( '-n', action='store_true', default=False, dest='n', help='not display the progressbar and colors' )
parser.add_argument( '-N', action='store_true', default=False, dest='N', help='also display number of logs' )
parser.add_argument( '-v', '--version', action='version', version='%(prog)s 2.0', help='Print program version' )	   
####

args = parser.parse_args()

hosts=[]

if (( args.lcm1 and args.lcm2 ) or args.condor ) : hosts = lcm1 + lcm2 
elif args.lcm1: hosts += lcm1
elif args.lcm2: hosts += lcm2
elif args.laur: hosts += laur
elif args.cuda: hosts += cuda
elif args.math: hosts += math
else: hosts = lcm1 + lcm2 + laur + cuda

if args.n:
   splitters = [('jake','LCM1'),('gin','LCM2'),('jacobi','CUDA'),('orion','LAUR')]
   progressbar = 0
else: 
   splitters = [('abe','\033[1;36mLCM1\033[0m'),('actarus','\033[1;32mLCM2\033[0m'),('jacobi','\033[1;35mCUDA\033[0m'),('eskimo','\033[1;33mLAUR\033[0m')]
   progressbar = 1
    
# Define port
port = 79 


#############################################
# Node class
# The bulk of the program is done here
#############################################

class Node(Thread):
    # Constructor
    def __init__(self,nome,port):
        # Fork the thread first thing
        Thread.__init__(self)
        # Variables initialization
        self.hostname = nome
        self.local = []
        self.remote = []
        self.lgod = []
        self.rgod = []
        # Former admins
        self.chuck = []
        # Node status
        self.up = True
        self.aval = True
        self.port = port
        self.timeout = False

    # 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")
        if "0 received" in ping.read():
            self.up = False
            return False
        else:
            return True
    
    # Open telnet connection
    def connect(self):
        # Try to connect first
        try: 
            self.conn = telnetlib.Telnet(self.hostname,self.port)
        # Inetd is down!
        except error, msg:
            self.aval = False
            return False
        return True
    
    # Read lcm_w ouput and fill accordingly the users-gods lists
    def read(self):    
        # Read data
        lista=''
        reachable=0
        try: lista = self.conn.read_until("EOF",1)
        except EOFError: reachable=1
        del self.conn
        
        if lista=='':
            if reachable == 0:
                self.timeout=True    

        # Split lines and fields
        righe = lista.splitlines()
        # Needed later for isempty
        self.users = len(righe)

        # Local-remote-god user check
        for x in righe:
            fields = x.split()

            if "tty" in fields[1]:
                if fields[0] in gods:
                    self.lgod.append(fields[0].lower())
                else:
                    if fields[0] not in chucknorris:
                        self.local.append(fields[0].lower())
            else:
                if fields[0] in gods:
                    if fields[0] not in self.lgod:
                        self.rgod.append(fields[0].lower())
                else:
                    if fields[0] not in self.local:
                        if fields[0] not in chucknorris:
                            self.remote.append(fields[0].lower())
            # Former admins
            if fields[0] in chucknorris:
                self.chuck.append(fields[0].lower())
    # Run method
    # This is the part that every thread executes in parallel
    # If host is up -> open connection -> if connection -> read data
    def run(self):
        if self.isup(): 
            if self.connect():
                self.read()
    
    # Is the host empty?
    def isempty(self):
        if (self.users > 0): 
            return False
        else:
            return True

    # Print output giorgio-style, host based (after threads have finished)
    def printlist(self):
        # Uniq equivalent and string conversion
        strlocal = ""
        strremote = ""
        strlgod = ""
        strrgod = ""
        strchuck = ""
        
    
        if not args.N:
            # A set cannot have duplicate entries: uniq equivalent
            for item in set(self.local):
                strlocal += str(item) + ' '
            for item in set(self.remote):
                strremote += str(item) +' '
            for item in set(self.lgod):
                strlgod += str(item) +' '
            for item in set(self.rgod):
                strrgod += str(item) +' '
            for item in set(self.chuck):
                strchuck += str(item) +' '
        else:
            # Print also how many times users are logged (added by jp)
            for item in set(self.local):
                strlocal += str(item) + '('+str(collections.Counter(self.local).values()[collections.Counter(self.local).keys().index(item)])+') '
            for item in set(self.remote):
                strremote += str(item) +'('+str(collections.Counter(self.remote).values()[collections.Counter(self.remote).keys().index(item)])+') '
            for item in set(self.lgod):
                strlgod += str(item) +'('+str(collections.Counter(self.lgod).values()[collections.Counter(self.lgod).keys().index(item)])+') '
            for item in set(self.rgod):
                strrgod += str(item) +'('+str(collections.Counter(self.rgod).values()[collections.Counter(self.rgod).keys().index(item)])+') '
            for item in set(self.chuck):
                strchuck += str(item) +'('+str(collections.Counter(self.chuck).values()[collections.Counter(self.chuck).keys().index(item)])+') '

        # Print a tag [C] aside CUDA hosts (jp)
        if self.hostname in cuda:
            self.hostname = self.hostname + ' '*(5-len(self.hostname))+'[C]'

        print ' ' + self.hostname + ' '*(15-len(self.hostname))+ red + strlgod + normal + pink + strrgod + normal + green + strlocal + normal + blue + strremote + normal+ turquoise+strchuck+normal


###################################################
# Main 
###################################################


# Get user groups
groups = os.popen("groups").readlines()[0].split()

# Setting colors
if progressbar > 0:
    red = "\033[1;31m"
    normal = "\033[0m"
    blue = "\033[1;34m"
    pink = "\033[1;35m"
    green = "\033[1;32m"
    turquoise ="\033[1;36m"
else:
    red = ""
    normal = ""
    blue = ""
    pink = ""
    green = ""
    turquoise =""


# Initialization of the four lists
downlist = ""
emptylist = ""
unavlist = ""
timeoutlist = ""

# Initialize the threadlist
threadlist = []

# Number of hosts (for progressbar)
num = len(hosts)

# Start the threads
for item in hosts:
    m=Node(item,port)
    threadlist.append(m)
    m.start()

# Used for progressbar
index = 0
print ' Querying '+str(num)+' hosts...'

# Rejoin them when their work is done
for thread in threadlist:
    thread.join()
    if progressbar > 0:
        # Progessbar
        index += 1
        stdout.write('\r ['+'='*index+'>'*(1-int(index/num))+' '*(num-index-1)+']')    
        stdout.flush()

if progressbar > 0:
    # Newline
    print '\n Done... ( %(t).3f s)' % {'t': (time() - start)}

# And now print!
for thread in threadlist:
    # lcm1-lcm2-laur splitters
    for pair in splitters:
        if thread.hostname in pair[0]: print '-' + pair[1] + normal +'----'
    
    # See what they are up to and print accordingly
    if thread.hostname in cuda:
        thread.hostname +='[C]'
    if args.math :
        if thread.hostname in math4:
            thread.hostname +='[M4]'
        if thread.hostname in math5:
                thread.hostname +='[M5]'
    if not thread.up:
        # Host down
        downlist += thread.hostname + ' '
    else:
        if not thread.aval:
            # Host unavailable
            unavlist += thread.hostname + ' '
            continue
        if thread.timeout:
            # Thread has timed out
            timeoutlist += thread.hostname + ' '
            continue
        if thread.isempty():
            # Host is empty
            emptylist += thread.hostname + ' '
        else:
            thread.printlist()
    
# Some final output
if progressbar > 0:
    print
    print 'Legenda: '+green + 'utente locale '+normal+blue+'utente remoto '+normal+red+'admin locale '+normal+pink+'admin remoto '+normal+turquoise+'nirvana'+normal
    print "         I nodi contrassegnati con [C] sono i nodi CUDA."
    if args.math:
        print "         I nodi contrassegnati con [M4] hanno Mathematica 4.0"
        print "         I nodi contrassegnati con [M5] hanno Mathematica 5.2"
    print

# Only in full mode 
#if ("-f" in str(params[0])) or ("-F" in str(params[0])) or ("-m" in str(params[0])) or ("-C" in str(params[0])):
if args.full or args.math or args.cuda: 
    print red+'Empty: '+normal+emptylist
    print red+'Timeout: '+normal+timeoutlist
    print red+'Unreachable: '+normal+downlist
    print red+'Unavailable: '+normal+unavlist

# Exit gracefully
exit(0)