#! /usr/bin/python
#
# This script checks that the user does not have any processes running
# on the node passed in as an argument if the user does not have any
# reservations for the node in the cluster queue.
#
# The intent of this script is to run it right after the ssh command
# that qlogin generates and then when that session ends it will warn
# the user that their processes may be killed.
#

import sys, re, os, pwd

commandsToIgnore = ["/usr/libexec/gam_server", "gconfd", "sshd:.*@notty", "/bin/ps auxwww"]

# Comand to get the processes running
# Setup the SGE environment:
os.environ["SGE_ROOT"]="/opt/gridengine"
os.environ["SGE_CELL"]="default"
os.environ["SGE_ARCH"]="lx26-amd64"
os.environ["SGE_EXECD_PORT"]="537"
os.environ["SGE_QMASTER_PORT"]="536"

qstatCommand = "/opt/gridengine/bin/lx26-amd64/qstat"
grepCommand = "/bin/grep"
sshCommand = "/usr/bin/ssh -x"
psCommand = "/bin/ps auxwww"

# Process the arguments
node = ""
port = 0
if (len(sys.argv) == 3):
    node = sys.argv[1]
    port = int(sys.argv[2])
else:
    print "Use: qlogin_logout_check.py compute-node.local PORT"
    sys.exit()

shortNode = ""
match = re.search("(compute-[0-9]-[0-9]+)\.local", node)
if (match):
    shortNode = match.group(1)
else:
    print "ERROR: failed to match host name."

username = pwd.getpwuid(os.getuid())[0]

# Find out if the user has any jobs scheduled on the node
child = os.popen(qstatCommand + " | " + grepCommand + " " + shortNode, "r")
for line in child.readlines():
    if (re.search(username, line)):
        sys.exit()
child.close()


printHeader = True

# Find out if the user has any processes on the node
command = sshCommand + " " + node + " \"" + psCommand + "\""
child = os.popen(command, "r")
for line in child.readlines():
    match = re.match("(\S+)\s+([0-9]+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(.*)", line)
    if (match):
        user = match.group(1)
        if (user != username):
            continue
        usercommand = match.group(11)
        pid = match.group(2)

        warnForCommand = True
        for ignore in commandsToIgnore:
            if (re.search(ignore, usercommand)):
                warnForCommand = False
                break
        if warnForCommand:
            if printHeader:
                printHeader = False
                print "\033[7;31mWARNING\033[0m: You have just released your queue reservation on " + shortNode
                print " but you still have processes running on this node."
                print " These processes \033[31;4mwill be killed\033[0m by the system at some point since you"
                print " do not have this node reserved."
                print " See \033[4mhttp://cva.stanford.edu/systems/cva_systems.html#InteractiveJobs\033[0m."
            print "\t(PID: " + pid +")\tcommand: " +usercommand
child.close()