#!/bin/bash
# Detect resources of a list of hosts and print OAR commands to create them in OAR database.

usage() {
        cat <<EOF
Usage:
    ${0##*/} [options] <hostfile>

Generate a script for the creation of the OAR resources of a list of hosts, by
connecting to each host and gathering system information.
Hosts are read one per line from hostfile (must be readable by user oar) or
STDIN (if hostfile is -).

Options:
  -H          handle HyperThreading
  -T          use the thread property for HyperThreading (imply -H)
  -o <file>   script filename
  -x          execute the script at the end and remove it
  -y          assume "yes" as answer to all prompts and run non-interactively
  -v          print the generated script
  -e <path>   path to the ssh command
  -n <path>   path to the oarnodesetting command
  -p <path>   path to the oarproperty command
  -c <path>   path to the oar.conf file
  -h          display this message

EOF

}

parse_opts() {
    while getopts "c:e:p:n:o:xyHTvh" opt; do
        case $opt in
        c)
            OARCONFFILE=$OPTARG
            ;;
        e)
            OPENSSH_CMD=$OPTARG
            ;;
        n)
            OARNODESETTING_CMD=$OPTARG
            ;;
        p)
            OARPROPERTY_CMD=$OPTARG
            ;;
        o)
            SCRIPT_FILE=$OPTARG
            ;;
        x)
            EXECUTE=1
            ;;
        v)
            PRINT=1
            ;;
        y)
            NON_INTERACTIVE=1
            ;;
        H)
            WITH_HT=1
            ;;
        T)
            WITH_HT=1
            WITH_THREAD_PROPERTY=1
            ;;
        h)
    				usage
            exit 0
            ;;
        \?)
            echo -e "Error: invalid option: -$OPTARG\n\n" >&2
    				usage 1>&2
            exit 1
          ;;
      esac
    done
}

OARCONFFILE="/etc/oar/oar.conf"
OPENSSH_CMD="/usr/bin/ssh"
OARPROPERTY_CMD="oarproperty"
OARNODESETTING_CMD="oarnodesetting"

parse_opts "$@"
shift $((OPTIND-1))
HOSTFILE=$1
shift 
OPTIND=1
parse_opts "$@"

if ! [ -r "$OARCONFFILE" ]; then
    echo -e "Error: cannot read OAR configuration file $OARCONFFILE\n\n" 1>&2
    usage 1>&2
    exit 1
fi

# check host file
if [ -z "$HOSTFILE" -o \( "$HOSTFILE" != "-" -a ! -r "$HOSTFILE" \) ]; then
    echo -e "Error: cannot read host file (warning: must be readable by the oar user)\n" 1>&2
    usage 1>&2
    exit 1
fi
. "$OARCONFFILE" || exit 20

if [ "$HOSTFILE" == "-" ]; then
    if [ -z "$NON_INTERACTIVE" ]; then
        echo -e "Error: reading nodes from STDIN requires the script to run non-interactive (-y option).\n" 1>&2
        usage 1>&2
        exit 1
    fi
    echo "# Reading nodes from STDIN..."
    exec 4<&0
    exec 0<&-
else
    echo "# Reading nodes from ${HOSTFILE}..."
    exec 4<$HOSTFILE
fi

if [ -z "$NON_INTERACTIVE" ]; then
    read -s -n 1 -p "# Is the OAR SSH key configured on all hosts? (press y to confirm)"
    echo
    [ "$REPLY" != "y" ] && exit 2
fi

# Check output file
if [ -n "$SCRIPT_FILE" ]; then
    if [ -e "$SCRIPT_FILE" ]; then
        if [ -z "$NON_INTERACTIVE" ]; then
            read -s -n 1 -p "# Script file $SCRIPT_FILE already exists, overwrite it? (press y to confirm)"
            echo
            [ "$REPLY" != "y" ] && exit 2
        fi
        echo -n > $SCRIPT_FILE || exit 4
    fi
else 
    SCRIPT_FILE=$(mktemp -p "" ${0##*/}.XXXXX)
fi

# Add properties
echo "$OARPROPERTY_CMD -a cpu || true" > $SCRIPT_FILE
echo "$OARPROPERTY_CMD -a core || true" >> $SCRIPT_FILE
[ -n "$WITH_HT" -a -n "$WITH_THREAD_PROPERTY" ] && echo "$OARPROPERTY_CMD -a thread || true" >> $SCRIPT_FILE
echo "$OARPROPERTY_CMD -c -a host || true" >> $SCRIPT_FILE
echo "$OARPROPERTY_CMD -a mem || true" >> $SCRIPT_FILE

# Init CPU id
CPU=$($OARNODESETTING_CMD --last-property-value cpu 2> /dev/null)
[ -n "$CPU" ] || CPU=0 
# Init CORE id
CORE=$($OARNODESETTING_CMD --last-property-value core 2> /dev/null)
[ -n "$CORE" ] || CORE=0
# Init THREAD id
if [ -n "$WITH_HT" -a -n "$WITH_THREAD_PROPERTY" ]; then
    THREAD=$($OARNODESETTING_CMD --last-property-value thread 2> /dev/null)
fi
[ -n "$THREAD" ] || THREAD=0

while read -u 4 host; do
    echo -n "#   Checking $host ... "

    # Store stderr and redirect it to /dev/null
    exec 3>&2
    exec 2> /dev/null

    mem=$($OPENSSH_CMD $host cat /proc/meminfo | perl -ane 'print "$F[1]" if /^MemTotal/')

    # Get CORE info in the form:
    # cpu:core
    # cpu:core
    # cpu:core
    # ...
    # of if WITH_TH:
    # cpu:core:thread
    # cpu:core:thread
    # cpu:core:thread
    # ...
    if [ -z "$WITH_HT" ]; then
        if [ "`$OPENSSH_CMD $host cat /proc/cpuinfo|grep 'physical id'`" != "" ]; then
            DATA=$($OPENSSH_CMD $host cat /proc/cpuinfo| perl -F'/\s*:\s*/' -ane 'chomp $F[1]; if ($F[0] eq "processor") { $processor = $F[1] } elsif ($F[0] eq "physical id") { $physical_id = $F[1]; $topo->{$physical_id}->{$processor} = 1 }; END{ foreach $physical_id (sort({$a <=> $b} keys(%$topo))) { foreach $processor (sort({$a<=>$b} keys(%{$topo->{$physical_id}}))) { print "$physical_id:$processor\n" } } }')
        else
            DATA=$($OPENSSH_CMD $host cat /proc/cpuinfo| perl -F'/\s*:\s*/' -ane 'chomp $F[1]; if ($F[0] eq "processor") { $topo->{$F[1]} = 1 }; END{ foreach $processor (sort ({$a <=> $b} keys(%$topo))) { print "$processor:$processor\n" } }')
        fi
    else
        DATA=$($OPENSSH_CMD $host cat /proc/cpuinfo| perl -F'/\s*:\s*/' -ane 'chomp $F[1]; if ($F[0] eq "processor") { $processor = $F[1] } elsif ($F[0] eq "physical id") { $physical_id = $F[1] } elsif ($F[0] eq "core id") { $core_id=$F[1]; $topo->{$physical_id}->{$core_id}->{$processor} = 1 }; END{ foreach $physical_id (sort({$a <=> $b} keys(%$topo))) { foreach $core_id (sort({$a <=> $b} keys (%{$topo->{$physical_id}}))) { foreach $processor (sort({$a<=>$b} keys(%{$topo->{$physical_id}->{$core_id}}))) { print "$physical_id:$core_id:$processor\n" } } } }')
    fi

    # Restore stderr
    exec 2>&3

    if [ -z "$mem" ] || [ -z "$DATA" ]; then
        echo "KO"
        cat 1>&2 <<EOF
Error: cannot retrieve info from node $host. Please check that the ssh connection as oar is ok, e.g.:
$ $OPENSSH_CMD $host cat /proc/cpuinfo
 
EOF
        exit 5
    fi

    ((mem = $mem / 1024 / 1024 + 1))
    if [ -z "$WITH_HT" ]; then
        prev_cpu="-1"
        for l in $DATA; do
            cpu=${l%%:*}
            core=${l##*:}
            [ "$cpu" != "$prev_cpu" ] && ((CPU++))
            ((CORE++))
            prev_cpu=$cpu
            echo "$OARNODESETTING_CMD -a -h $host -p host=$host -p cpu=$CPU -p core=$CORE -p cpuset=$core -p mem=$mem" >> $SCRIPT_FILE
        done
    else
        prev_cpu="-1"
        prev_core="-1"
        tmp_cpuset=""
        cpu=""
        core=""
        for l in $DATA; do
            cpu=${l%%:*}
            ll=${l#*:}
            core=${ll%%:*}
            thread=${l##*:}
            if [ -z "$WITH_THREAD_PROPERTY" ]; then
                if [ "$core" != "$prev_core" ]; then
                    if [ -n "$tmp_cpuset" ]; then
                        echo "$OARNODESETTING_CMD -a -h $host -p host=$host -p cpu=$CPU -p core=$CORE -p cpuset=$tmp_cpuset -p mem=$mem" >> $SCRIPT_FILE
                    fi
                    tmp_cpuset="$thread"
                else
                    tmp_cpuset="$tmp_cpuset,$thread"
                fi
            fi
            [ "$cpu" != "$prev_cpu" ] && ((CPU++))
            [ "$core" != "$prev_core" ] && ((CORE++))
            ((THREAD++))
            if [ -n "$WITH_THREAD_PROPERTY" ]; then
                echo "$OARNODESETTING_CMD -a -h $host -p host=$host -p cpu=$CPU -p core=$CORE -p thread=$THREAD -p cpuset=$thread -p mem=$mem" >> $SCRIPT_FILE
            fi
            prev_cpu=$cpu
            prev_core=$core
        done
        if [ -z "$WITH_THREAD_PROPERTY" ]; then
            echo "$OARNODESETTING_CMD -a -h $host -p host=$host -p cpu=$CPU -p core=$CORE -p cpuset=$tmp_cpuset -p mem=$mem" >> $SCRIPT_FILE
        fi
    fi
    echo "OK"
done
echo -ne "# Done\n\n"

if [ -n "$PRINT" ]; then
    echo -ne "# Generated script content:\n\n"
    cat $SCRIPT_FILE
    echo
fi
if [ -n "$EXECUTE" ]; then
    echo -ne "# Executing script...\n\n"
    . $SCRIPT_FILE
    echo -ne "\n# Removing script file.\n"
    rm $SCRIPT_FILE
else
    cat <<EOF
# Please please find the generated script here: $SCRIPT_FILE
# It everything looks correct to you, source it to create the resources, by running:
# . $SCRIPT_FILE

EOF
fi
