Tool for viewing lab machine availability and active sessions.

cs-avail.sh

#!/bin/bash

AUTO_LOGIN="true"
LIST_ALL="false"
VERBOSE="false"
USERNAME="${USER}"
TARGETNAME="xinu"
SEARCH_NAME=""

MIN_USER_COUNT=-1
MIN_USER_ID=""
SEARCH_HITS=0
PREV_SSH_SUCCESS="false"

function help-text {
    echo "cs-avail"
    echo "    Assess the number of users logged into a set of lab machines and"
    echo "    connect to the best available."
    echo
    echo "Usage:"
    echo "    cs-avail [OPTIONS] [-x prefix] [-u username] [-f search]"
    echo
    echo "    cs-avail                Default configuration (auto-login for xinu)"
    echo "    cs-avail -h             Display this message"
    echo "    cs-avail -v             Verbose -- display the user list"
    echo "    cs-avail -l             List user stats for all machines"
    echo "    cs-avail -b             Disable auto-login"
    echo "    cs-avail -x moore       Connect to machines with the prefix 'moore'. Default: 'xinu'"
    echo "    cs-avail -f friend1     Search for active sessions w/ username 'friend1'"
    echo "    cs-avail -u test        Connect using test@host instead of ${USER}@host"
    echo
    echo "IMPORTANT:"
    echo "    This script assumes you have an SSH key setup for the targeted machines."
    echo "    Without a key, you'll be required to type your password for each connection."
    echo "    To easily create and deploy keys, see tiny.cc/install-key"
    echo
}

while getopts ":hvlbx:u:f:" opt; do
    case ${opt} in
        h )
            help-text
            exit 0
            ;;
        v )
            VERBOSE="true"
            ;;
        l )
            LIST_ALL="true"
            ;;
        b )
            AUTO_LOGIN="false"
            ;;
        x )
            TARGETNAME="$OPTARG"
            ;;
        u )
            USERNAME="$OPTARG"
            ;;
        f)
            SEARCH_NAME="$OPTARG"
            ;;
        \? )
            echo "Unknown option: -$OPTARG" 1>&2
            help-text
            exit 1
            ;;
    esac
done

for ((i=1;1;i++)); do
    
    # Prefix a zero if we need to
    id="${i}"
    if [ ${i} -lt 10 ]; then
        id="0${i}"
    fi
    
    # Prevent SSH attempt to ourself
    # Compare hostname up to first .
    if [ "${TARGETNAME}${id}" == "${HOSTNAME%%.*}" ]; then
        USER_LIST="$(users)"
        SSH_EXIT_CODE=0
    else
        USER_LIST="$(ssh ${USERNAME}@${TARGETNAME}${id}.cs.purdue.edu 'users' 2>&1)"
        SSH_EXIT_CODE=${?}
    fi
    
    # Check for connection issues
    # Deduce when we've hit the last numbered machine
    if [ ${SSH_EXIT_CODE} -ne 0 ] && [ "${PREV_SSH_SUCCESS}" == "false" ]; then
        echo
        echo "SSH connection to ${TARGETNAME}${id} failed."
        exit
    elif [ ${SSH_EXIT_CODE} -ne 0 ]; then
        # No more
        break
    else
        echo "Connected to ${TARGETNAME}${id}..."
    
        PREV_SSH_SUCCESS="true"
        
        # Search for the targeted username
        if [ ! -z "${SEARCH_NAME}" ]; then
            result=$(echo "${USER_LIST}" | grep --color=always -w "${SEARCH_NAME}")
            
            if [ ${?} -eq 0 ]; then
                echo
                echo "Found ${SEARCH_NAME} on ${TARGETNAME}${id}."
                
                if [ "${VERBOSE}" == "true" ]; then
                    echo "${result}"
                fi
                
                echo
                
                SEARCH_HITS=$((SEARCH_HITS+1))
            fi
            
            # Don't do anything else
            continue
        fi
        
        # Do user count
        USER_COUNT=$(echo "${USER_LIST}" | wc -w)
        echo "${USER_COUNT} active sessions."
        
        if [ "${VERBOSE}" == "true" ]; then
            echo "${USER_LIST}"  
        fi

        if [ ${MIN_USER_COUNT} -eq -1 ] || [ ${USER_COUNT} -lt ${MIN_USER_COUNT} ]; then
            MIN_USER_COUNT=${USER_COUNT}
            MIN_USER_ID=${id}
        fi
    fi
    
    echo
    
    # Don't waste time checking more if we already found an empty one
    if [ "${LIST_ALL}" == "false" ] && [ ${MIN_USER_COUNT} -eq 0 ]; then
        break
    fi

done

# Handle search results
if [ ${SEARCH_HITS} -gt 0 ]; then
    echo
    echo "Found ${SEARCH_HITS} session(s) associated with ${SEARCH_NAME}."
    echo
    exit 0
elif [ ! -z "${SEARCH_NAME}" ]; then
    echo
    echo "Could not find session(s) associated with ${SEARCH_NAME}."
    echo
    exit 1
fi

if [ "${LIST_ALL}" == "false" ] && [ "${AUTO_LOGIN}" == "true" ]; then

    echo "Starting session on best available host (${TARGETNAME}${MIN_USER_ID})..."
    echo

    # Start the ssh session
    ssh ${USERNAME}@${TARGETNAME}${MIN_USER_ID}.cs.purdue.edu
    exit 0
    
elif [ "${AUTO_LOGIN}" == "false" ]; then
    
    echo
    echo "Best/first available host is ${TARGETNAME}${MIN_USER_ID}."
    echo

fi