whoall/whoall

#!/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
#
# Last modified: March 2016 - blue


from time import time
start = time()

import argparse
import os
import telnetlib
from collections import Counter
from threading import Thread
from socket import error
from sys import exit,stdout


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

# Define 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.1', help='Print program version')

# Parse arguments
args = parser.parse_args()

# Define port on which the lcm_w service is listening
# NB at the time of writing lcm_w is managed by inetd
PORT = 79


#############################################
# Node class
# Each node executes a `who` query in a separate thread
#############################################

class Node(Thread):
    # Constructor
    def __init__(self, name, location, math4=False, math5=False,
                 cuda=False, condor=False, port=PORT):
        # Fork the thread first thing
        Thread.__init__(self)
        # Variables initialization
        self.hostname = name
        self.location = location
        if math4:
            self.math_version = 'M4'
        elif math5:
            self.math_version = 'M5'
        else:
            self.math_version = 'NA'
        self.cuda = cuda
        self.condor = condor
        self.local = []
        self.remote = []
        self.lgod = []
        self.rgod = []
        # Former admins
        self.chuck = []
        # Node status
        self.up = True
        self.avail = 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:
            self.avail = 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 hostname (plus qualifiers if needed)
    def qualifiedname(self):
        tmp = self.hostname

        # Add math version to hostname if required
        if args.math:
            tmp += '[' + self.math_version + ']'

        # Add cuda tag [C] to hostname for CUDA hosts (jp)
        # FIXME I don't think we need this. Cuda nodes have their own section
        # if self.cuda:
        #     self.hostname += ' '*(5-len(self.hostname))+'[C]'

        return tmp


    # Print output giorgio-style, host based (after threads have finished)
    def printlist(self):
        # Uniq equivalent and string conversion
        strlocal = ""
        strremote = ""
        strlgod = ""
        strrgod = ""
        strchuck = ""

        if args.N:
            # Also print how many times users are logged (jp+blue)
            lcounter = Counter(self.local)
            for user in lcounter:
                strlocal += user + '(' + str(lcounter[user]) + ')'
            rcounter = Counter(self.remote)
            for user in rcounter:
                strremote += user + '(' + str(rcounter[user]) + ') '
            lgcounter = Counter(self.lgod)
            for user in lgcounter:
                strlgod += user + '(' + str(lgcounter[user]) + ') '
            rgcounter = Counter(self.rgod)
            for user in rgcounter:
                strrgod += user + '(' + str(rgcounter[user]) + ') '
            ccounter = Counter(self.chuck)
            for user in ccounter:
                strchuck += user + '(' + str(ccounter[user]) + ') '
        else:
            # 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) +' '

        # Print out hostname and connected users
        print ' ' + self.qualifiedname() + ' '*(15-len(self.qualifiedname())) \
                + red + strlgod + normal + pink + strrgod + normal + green \
                + strlocal + normal + blue + strremote + normal \
                + turquoise + strchuck + normal
### end class Node


### USER GROUPS ###
# Current admins/150
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']


### HOST LIST ###
# Only edit here to add/remove/change hostlist

nodes = [
    Node('abe',     'LCM1',     math5=True,     condor=True),
    Node('crash',   'LCM1',     math4=True,     condor=True),
    Node('duke',    'LCM1',     math4=True,     condor=True),
    Node('glados',  'LCM1',     math5=True,     condor=True),
    Node('lara',    'LCM1',                     condor=True),
    Node('link',    'LCM1',     math5=True,     condor=True),
    Node('king',    'LCM1',     math5=True,     condor=True),
    Node('pang',    'LCM1',     math5=True,     condor=True),
    Node('pong',    'LCM1',     math5=True,     condor=True),
    Node('snake',   'LCM1',     math5=True,     condor=True),
    Node('sonic',   'LCM1',     math4=True,     condor=True),
    Node('spyro',   'LCM1',     math5=True,     condor=True),
    Node('yoshi',   'LCM1',     math5=True,     condor=True),
    Node('actarus', 'LCM2',     math4=True,     condor=True),
    Node('elwood',  'LCM2',     math4=True,     condor=True),
    Node('gex',     'LCM2',                     condor=True),
    Node('gin',     'LCM2',     math4=True,     condor=True),
    Node('jake',    'LCM2',     math4=True,     condor=True),
    Node('kirk',    'LCM2',     math4=True,     condor=True),
    Node('martini', 'LCM2',     math4=True,     condor=True),
    Node('picard',  'LCM2',     math4=True,     condor=True),
    Node('q',       'LCM2',     math4=True,     condor=True),
    Node('raziel',  'LCM2',     math4=True,     condor=True),
    Node('sarek',   'LCM2',     math4=True,     condor=True),
    Node('spock',   'LCM2',                     condor=True),
    Node('tron',    'LCM2',     math4=True,     condor=True),
    Node('worf',    'LCM2',     math4=True,     condor=True),
    Node('zombie',  'LCM2',     math4=True,     condor=True),
    Node('eskimo',  'LAUR'),
    Node('orion',   'LAUR'),
    Node('tilde',   'LAUR'),
    Node('jacobi',  'CUDA', cuda=True),
    Node('yukawa',  'CUDA', cuda=True),
    Node('tesla',   'CUDA', cuda=True),
]


### SECTION SPLITTERS ###
if args.n:
   splitters = { 'LCM1': '-LCM1----',
                 'LCM2': '-LCM2----',
                 'CUDA': '-CUDA----',
                 'LAUR': '-LAUR----' }
else:
   splitters = { 'LCM1': '-\033[1;36mLCM1\033[0m----',
                 'LCM2': '-\033[1;32mLCM2\033[0m----',
                 'CUDA': '-\033[1;35mCUDA\033[0m----',
                 'LAUR': '-\033[1;33mLAUR\033[0m----' }


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

# Set terminal colors
if not args.n:
    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 =""

# Check whether we want a progressbar displayed
progressbar = not args.n

# Filter hostlist according to arguments
if args.lcm1:
    nodes = [ node for node in nodes if node.location == 'LCM1' ]
if args.lcm2:
    nodes = [ node for node in nodes if node.location == 'LCM2' ]
if args.laur:
    nodes = [ node for node in nodes if node.location == 'LAUR' ]
if args.cuda:
    nodes = [ node for node in nodes if node.location == 'CUDA' ]
if args.math:
    nodes = [ node for node in nodes if node.math_version != 'NA' ]

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

# Start the threads
for node in nodes:
    node.start()

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

# Try to join each node in turn
# Remove their reference from nodescopy as they finish
iterator = 0
nodescopy = nodes[:] # shallow copy of nodes list
while len(nodescopy) > 0:
    node = nodescopy[iterator%len(nodescopy)]
    iterator += 1
    node.join(0)
    if node.isAlive():
        # Not done yet, let's come back later
        continue
    else:
        # Node is done, remove it from nodescopy
        nodescopy.remove(node)
        # Update progressbar
        if progressbar:
            # Progessbar
            index += 1
            stdout.write('\r ['
                         + '='*index
                         + '>'*(1-int(index/num))
                         + ' '*(num-index-1) + ']')
            stdout.flush()

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

if args.full or args.math or args.cuda:
# Save down/unreachable/unavailable nodes to separate lists
    downlist = ''
    timeoutlist = ''
    unavlist = ''
    emptylist = ''
    for node in nodes:
        if not node.up:
            downlist += node.qualifiedname() + ' '
        elif node.timeout:
            timeoutlist += node.qualifiedname() + ' '
        elif not node.avail:
            unavlist += node.qualifiedname() + ' '
        elif node.isempty():
            emptylist += node.qualifiedname() + ' '

nodes = [ node for node in nodes if node.up
                                    and not node.timeout
                                    and node.avail
                                    and not node.isempty() ]

# Split notes by location
lcm1nodes = [ node for node in nodes if node.location == 'LCM1' ]
lcm2nodes = [ node for node in nodes if node.location == 'LCM2' ]
laurnodes = [ node for node in nodes if node.location == 'LAUR' ]
cudanodes = [ node for node in nodes if node.location == 'CUDA' ]

# Print out
if len(lcm1nodes):
    print splitters['LCM1']
    for node in lcm1nodes:
        node.printlist()
if len(lcm2nodes):
    print splitters['LCM2']
    for node in lcm2nodes:
        node.printlist()
if len(laurnodes):
    print splitters['LAUR']
    for node in laurnodes:
        node.printlist()
if len(cudanodes):
    print splitters['CUDA']
    for node in cudanodes:
        node.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
    # FIXME I don't think we need this. Cuda nodes have their own section
    # 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 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)