#!/usr/bin/python
# - *- coding: utf- 8 - *-
#
# 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

# TODO
# - improve means of separation between lcm1, lcm2, laur
# - improve handling of options (e.g. "whoall -1 -m" should return
#   nodes in lcm1 which have mathematica installed)

from time import time
start = time()

# Import required 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 list 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', dest='lcm1',
                    help='only display LCM1 hosts ')
parser.add_argument('-2', '--lcm2', action='store_true', dest='lcm2',
                    help='only display LCM2 hosts ')
parser.add_argument('-l', '--laur', action='store_true', dest='laur',
                    help='only display LAUR hosts ')
parser.add_argument('-c', '--condor', action='store_true', dest='condor',
                    help='only display condor nodes ')
parser.add_argument('-C', '--cuda', action='store_true', dest='cuda',
                    help='only display CUDA hosts ')
parser.add_argument('-f', '--full', action='store_true', dest='full',
                    help='display all nodes including empty, \
                    unreachable and unavailable')
parser.add_argument('-m', '--math', action='store_true', dest='math',
                    help='only Mathematica hosts ')
parser.add_argument('-n', action='store_true', dest='n',
                    help='not display the progressbar and colors')
parser.add_argument('-N', action='store_true', 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

# FIXME it would be better not to rely on host names for splitting
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 users-gods lists accordingly
    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 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:
            # Also print how many times users are logged (jp+blue)
            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] alongside 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)