#! /usr/bin/python
# 
# This script is intended to prevent users from reserving more than one
# interactive session on the server. It does this by looking to see if the
# user already has a qlogin session in the queue, and if so it ssh's into
# that node. If the user does not have one then qlogin is called to create
# one. If it is called with "-new" a new session is requested.
#
# 02-mar-06 - davidbbs - initial version
# 11-apr-06 - davidbbs - modified to pass arguments to qlogin
# 12-may-06 - davidbbs - modified to check that it is running on the head node only
#

import sys, re, os, socket


DEBUG = False
headNodeName = "bagels.stanford.edu"
qloginCommand = "/opt/gridengine/bin/lx26-amd64/qlogin"
sshCommand = "/usr/bin/ssh"
qstatCommand = "/opt/gridengine/bin/lx26-amd64/qstat | /bin/grep QLOGIN"

if (socket.gethostname() != headNodeName):
	print "ERROR: Please run qlogin from the head node ("+headNodeName+") not this node (" + socket.gethostname() +")."
	sys.exit()

newSession = False

# Check if the user wants to force a new qlogin session
argStartIndex = 1
if (len(sys.argv) > 1):
    if (sys.argv[1] == "-new"):
        print "Forcing new interactive queue reservation..."
        newSession = True
	argStartIndex = 2
	
# Check if there are any commands to pass in to qlogin, and if so add them to the command
if (len(sys.argv) > argStartIndex):
	for i in range(argStartIndex,len(sys.argv)):
		qloginCommand += " " + sys.argv[i]
if (DEBUG):
	print qloginCommand

user = os.environ["USER"]

if (newSession == False):
    newSession = True
    # Find out if the user has a qlogin session
    child = os.popen(qstatCommand, "r")
    # For each line in the output extract the jobID, the user, and the machine
    for line in child.readlines():
        wMatch = re.search("\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)", line)
        if (wMatch):
            # Extract the fields
            qloginUser = wMatch.group(4)
            machine = wMatch.group(8)
            # Extract the machine in the format compute-x-y.local
            computeMatch = re.search("(\S+).q@(\S+)\.", machine)
            machineName = "c0-0.local"
            if (computeMatch):
                machineName = computeMatch.group(2) + ".local"

            if (qloginUser == user):
                print ""
                print "\033[31mYou already have an interactive session reserved on " + machineName + ".\033[0m"
                print "  If you want to reserve another session run \"qlogin -new\"."
                print "  See \033[4mhttp://cva.stanford.edu/systems/cva_systems.html#InteractiveJobs\033[0m"
                print "\033[31mOpening ssh connection to previously reserved compute node " + machineName + "...\033[0m\n"
                os.system(sshCommand + " " + machineName)
                newSession = False
                sys.exit()
    child.close();

if (newSession):
    print ""
    print "\033[31mReserving a new interactive session...\033[0m"
    print ""
    os.system(qloginCommand)
    # We should really get the node to which we connected here and
    # check if the user has any other qlogin sessions to that node.
    # If not we should kill all sshd processes owned by that user on the node.
    # This is actually really hard to do since we can't get the host from
    # this command. (I.e., if we try to read the error stream we stop the
    # whole qlogin session.)
    #
    # To solve this we call teh qlogin_logout_check.py script from the
    # rocks-qlogin.sh scrip (which is what calls ssh) when it exits.)
    #
    #print "******************************************************"
    #print "You have released your interactive queue reservation."
    #print " Any other jobs or terminals you have on the machine"
    #print "  you just exited may be terminated automatically."
    #print "******************************************************"