sssh

#!/bin/bash

PROGRAM_NAME="sssh"
VERSION_STRING="0.3"

# If the user does not have $HOME/.scalr/config.ini,
# DEFAULT_SCALR_CONFIG will be copied to $HOME/.scalr/config.ini, for
# the convinience.
DEFAULT_SCALR_CONFIG=${DEFAULT_SCALR_CONFIG:="/usr/local/etc/scalr/config.ini"}

# Each farm has a role with a VM named 'router-...' and it will have
# an external IP.  This VM will act as a ssh proxy server to connect
# other VM if the other VM does not have an external IP.  To override
# this, define SSH_PROXY_ENDPOINT as below.

# To connect a server without external IP, you'll need a SSH routing
# server for the connection.  This routing server should be accessible
# from the client machine that will run sssh(1), and should be able to
# connect the target server.  In addition to this, this routing server
# need nc(1) command in PATH.
SSH_PROXY_ENDPOINT=${SSH_PROXY_ENDPOINT:="root@www.xxx.yyy.zzz"}

# To connect a server in farm (farm-id = 409), the PEM file (409.pem)
# should be either in $PEM_DIR/409.pem or $HOME/.ssh/409.pem
PEM_DIR=/usr/local/etc/sssh/pem

nocache=""
farmid=""
nocolor=""

if [ ! -t 1 ]; then
    export nocolor=1
fi

function tp() {
    if [ -n "$nocolor" ]; then
        return 0;
    fi
    if [ -z "$TERM" ]; then
        return 1;
    fi
    tput "$@" 2>/dev/null
    return $?
}

function error() {
    local status

    status="$1"
    shift

    echo "$PROGRAM_NAME: $@" 1>&2
    if [ 0 -ne "$status" ]; then
        exit "$status"
    fi
}

function farm-id() {
    local fid
    if [ -n "$farmid" ]; then
        echo "$farmid"
    else
        if [ -r $HOME/.sssh/farmid ]; then
            read fid < $HOME/.sssh/farmid
            echo "$fid"
        else
            error 1 "no saved farm id.  Use set-farm command"
        fi
    fi
}

function identity-file() {
    # Usage: identity-file FID
    local fid=$1
    if [ -r $HOME/.ssh/$fid.pem ]; then
        echo "$HOME/.ssh/$fid.pem"
    elif [ -r $PEM_DIR/$fid.pem ]; then
        echo "$PEM_DIR/$fid.pem"
    else
        return 1
    fi
    return 0
}

function valid-cache-p() {
    # Usage: valid-cache-p cache-file expiration_in_seconds
    local created expiration now

    if [ -n "$nocache" ]; then
        return 1
    fi

    created=$(stat -c %Y "$1" 2>/dev/null)
    now=$(date +%s 2>/dev/null)
    expiration=${2:-300}        # Default: 300 seconds

    if [ -z "$created" -o -z "$now" ]; then
        return 1;
    fi

    if [ $(stat -c %s "$1") -eq 0 ]; then
    	return 1;
    fi

    if [ "$expiration" -ge $((now - created)) ]; then
        return 0;
    fi
    return 1;
}


function find-bastion-server() {
    local bfid=$(find-bastion-farm)

    if [ -z "$bfid" ]; then
        return 1
    fi

    scalr list-servers -f "$bfid" -c ExternalIP,Status | \
        awk -F '|' \
            -f <(cat - <<-'EOD'
NF == 4 {
	sub(/^ */, "", $2);	sub(/ *$/, "", $2);
	sub(/^ */, "", $3);	sub(/ *$/, "", $3);

	if ($3 == "Running") {
		print $2
		exit 0
	}
}
END {
	exit 1
}
EOD
                )
    
}

function find-bastion-farm() {
    local current

    rm -f $HOME/.sssh/bfid
    
    scalr list-farms | \
        awk -F '|' \
            -v "bfid=$HOME/.sssh/bfid" \
            -f <(cat - <<-'EOD'
NF == 6 { 
	sub(/^ */, "", $2);	sub(/ *$/, "", $2);
	sub(/^ */, "", $3);	sub(/ *$/, "", $3);
	sub(/^ */, "", $4);	sub(/ *$/, "", $4);
	sub(/^ */, "", $5);	sub(/ *$/, "", $5);

	if ($3 ~ /bastion$/) {
		print $2
		print $2 >bfid
		exit 0
	}
}
END {
	exit 1
}
EOD
)
}

function find-router() {
    local fid cache bastion

    fid=$(farm-id)
    if [ -z "$fid" ]; then
        exit 1
    fi

    bastion=$(find-bastion-server)
    if [ -n "$bastion" ]; then
        echo "$bastion"
        return 0
    fi
    
    scalr list-servers -f "$fid" -c ExternalIP,Name | \
    awk -F '|' \
        -f <(cat - <<-'EOD'
BEGIN { id = 0; c_n="" }
NF == 4 { 
	sub(/^ */, "", $2);	sub(/ *$/, "", $2);
	sub(/^ */, "", $3);	sub(/ *$/, "", $3);

	if ($3 ~ /^router/ && $2 != "None") {
		print $2
		exit 0
	}
}
END { exit 1 }
EOD
)
}


function list-servers() {
    local fid cache

    fid=$(farm-id)
    if [ -z "$fid" ]; then
        exit 1
    fi

    cache="$HOME/.sssh/$fid"

    if ! valid-cache-p "$cache"; then
        scalr list-servers -f "$fid" -c ExternalIP,InternalIP,Name > "$cache"
    fi
        

    cat "$cache" | \
    awk -F '|' \
        -v "c_bo=$(tp bold)" \
        -v "c_r=$(tp sgr0)" \
        -v "c_bl=$(tp setaf 4)" \
        -v "c_gr=$(tp setaf 2)" \
        -f <(cat - <<-'EOD'
BEGIN { id = 0; c_n="" }
NF == 5 { 
	sub(/^ */, "", $2);	sub(/ *$/, "", $2);
	sub(/^ */, "", $3);	sub(/ *$/, "", $3);
	sub(/^ */, "", $4);	sub(/ *$/, "", $4);

	if ($2 != "ExternalIP") {
		if ($2 == "None")
			printf "  [%s%2d%s]  %s%15s%s %s%15s%s  %s\n", \
				c_bo, id, c_r, c_n, $2, c_n, c_gr, $3, c_r, $4;
		else
			printf "  [%s%2d%s]  %s%s%15s%s %s%15s%s  %s\n", \
				c_bo, id, c_r, c_bo, c_bl, $2, c_r, c_gr, $3, c_r, $4;
		id += 1
	}
}
EOD
)
}

function connect-server() {
    local given fid pem eip iip cache proxy proxyport
    local OPTIND OPTARG
    local -a CMD

    CMD+=("ssh")
    proxyport=22
    while getopts 1246AaCfgKkMNnqsTtVvXxYyb:c:D:E:e:F:I:i:L:l:m:O:o:p:P:Q:R:S:W:w: opt; do
        case $opt in
            [bcDEeFIiLlmOopQRSWw])
                #echo "opt($opt) with arg($OPTARG)"
                CMD+=("-$opt" "$OPTARG")
                ;;
            P)
                proxyport="$OPTARG"
                ;;
            \?)
                error 1 "invalid option -- $opt"
                ;;
            *)
                #echo "opt($opt)"
                CMD+=("-$opt")
                ;;
        esac
    done

    shift $((OPTIND - 1))
    given="$1"
    fid=$(farm-id)
    if [ -z "$fid" ]; then
        error 1 "no farm is selected, you may forget to run set-farm command"
    fi
    if [ -z "$given" ]; then
        error 1 "argument(ID) required"
    fi
    shift

    cache="$HOME/.sssh/$fid"
    if ! valid-cache-p "$cache"; then
        scalr list-servers -f "$fid" -c ExternalIP,InternalIP,Name > "$cache"
    fi

    cinfo=/tmp/sssh.$$
    trap "rm -f $cinfo; exit 1" SIGINT SIGHUP SIGTERM

    cat "$cache" | \
    awk -F '|' -f <(cat - <<-EOD
BEGIN { id = 0 }
NF == 5 {
	#print \$0
	sub(/^ */, "", \$2);	sub(/ *$/, "", \$2);
	sub(/^ */, "", \$3);	sub(/ *$/, "", \$3);
	sub(/^ */, "", \$4);	sub(/ *$/, "", \$4);

	if (\$2 != "ExternalIP") {
		if (id == $given) {
			if (\$2 == "None")
				host=""
			else
				host=\$2
			printf "%s#%s\n", host, \$3
			exit 0
		}
		id += 1
	}
}
EOD
) > $cinfo
    eip=$(cat $cinfo | awk -F $'#' '{ print $1 }')
    iip=$(cat $cinfo | awk -F $'#' '{ print $2 }')

    rm -f "$cinfo"
    trap "" SIGINT SIGHUP SIGTERM

    pem=$(identity-file $fid)
    if [ -z "$pem" ]; then
        error 1 "You don't have the pem file ($pem)"
        exit 1
    fi
    bpem=$(identity-file $(cat $HOME/.sssh/bfid 2>/dev/null) 2>/dev/null)
    #echo "bpem: $bpem"
    if [ -z "$eip" -o -n "$bpem" ]; then
        if ! read proxy < $HOME/.sssh/router; then
            proxy=$SSH_PROXY_ENDPOINT
        fi

        if [ -z "$bpem" ]; then
            bpem="$pem"
        fi
           
        CMD+=(-i "$pem" -o "ProxyCommand ssh -p $proxyport -i $bpem $proxy nc %h %p" root@${iip} "$@")
    else
        CMD+=(-i "$pem" root@${eip} "$@")
    fi

    #echo "cmd: ${CMD[*]}"
    "${CMD[@]}"
}

function list-farms() {
    local current

    current=$(cat ~/.sssh/farmid)
    scalr list-farms | \
    awk -F '|' \
        -v "c_bo=$(tp bold)" \
        -v "c_r=$(tp sgr0)" \
        -v "c_bl=$(tp setaf 4)" \
        -v "c_gr=$(tp setaf 2)" \
        -v "c=$current" \
        -f <(cat - <<-'EOD'
NF == 6 { 
	sub(/^ */, "", $2);	sub(/ *$/, "", $2);
	sub(/^ */, "", $3);	sub(/ *$/, "", $3);
	sub(/^ */, "", $4);	sub(/ *$/, "", $4);
	sub(/^ */, "", $5);	sub(/ *$/, "", $5);

        if ($2 == c)
		mark = "*"
	else
		mark = " "

	if ($5 == "Running") {
		printf "  [%s%4d%s]  %s%s%32s%s%s (%s%s%s%s):%*s%s\n", \
			c_bo, $2, c_r, c_bo, c_bl, $3, c_r, mark, c_bo, c_gr, $5, c_r, 10 - length($5), " ", $4;
	}
	else if ($5 != "Status") {
		if ($2 == "")
			printf "   %3s   %32s  %s  %*s%s\n", "", "", "", 10, " ", $4
		else
			printf "  [%4d]  %32s%s (%s):%*s%s\n", $2, $3, mark, $5, 10 - length($5), " ", $4;
	}
}
EOD
)
}


function set-farm() {
    if [ "$#" -ne 1 ]; then
        error 1 "Usage: set-farm FARM-ID"
    fi

    rm -f "$HOME/.sssh/farmid" "$HOME/.sssh/router" "$HOME/.sssh/bfid"

    echo "$1" > $HOME/.sssh/farmid

    echo "router: $(find-router)"
    echo "root@$(find-router)" > $HOME/.sssh/router
}


function list-environments() {
    local current
    current=$(grep ^env_id ~/.scalr/config.ini 2>/dev/null | \
        grep -o "[0-9][0-9]*$")
    scalr list-environments | \
    awk -F '|' \
        -v "c_bo=$(tp bold)" \
        -v "c_r=$(tp sgr0)" \
        -v "c_bl=$(tp setaf 4)" \
        -v "c_gr=$(tp setaf 2)" \
        -v "c=$current" \
        -f <(cat - <<-'EOD'
NF == 4 { 
	sub(/^ */, "", $2);	sub(/ *$/, "", $2);
	sub(/^ */, "", $3);	sub(/ *$/, "", $3);

	if ($2 != "ID")
		printf "  [%s%3d%s]  %s%s%s%s %s\n", c_bo, $2, c_r, c_bo, c_bl, $3, c_r, (c == $2) ? "*" : ""
}
EOD
)
}


function set-environment() {
    local envid
    if [ "$#" -ne 1 ]; then
        error 1 "Usage: set-env ENVIRONMENT-ID"
    fi
    envid="$1"
    if [ ! -f "$HOME/.scalr/config.ini" ]; then
        error 0 "missing scalr configuration"
        error 1 "you need to run 'scalr configure'"
    fi

    sed -i -e "s/^env_id *=.*$/env_id = ${envid}/" "$HOME/.scalr/config.ini"
}


function version-and-exit() {
    echo "$PROGRAM_NAME version $VERSION_STRING"
    exit 0
}

function help-and-exit() {
    cat <<EOF
$PROGRAM_NAME: list|connect Scalr-managed server.
Usage: $PROGRAM_NAME [OPTION...] COMMAND [COMMAND-ARGUMENT...]

Options:

    -f FID   Use FID-th farm instead of saved one.

    -n       Ignore any cached data.

    -h       show help and exit
    -v       print version string and exit


Commands:

    list                                     List servers
    connect [SSH-OPTIONS...] ID [COMMAND]    Connect to the server(ID)

    env                                      List environments
    set-env ID                               Set current environment to env(ID)

    farms                                    List farms
    set-farm ID                              Set current farm to the farm(ID)


Example:

    First, you may need to list the available environment using 'env' command,
    then, set the default environment using 'set-env' command.

    Next, you may need to list the available farm using 'farms' command,
    then, set the default farm using 'set-farm' command.

    Finally, you may need to list the available servers using 'list' command,
    then, connect to the server using 'connect' command.   'connect' can accept
    SSH options if needed.  For example, if you want to connect 3rd VM using
    non-standard port (e.g. 443), you may use 'sssh connect -p 443 3'.  If you
    just want to run specific command, you may pass additional shell command 
    like 'sssh connect ls /'

    You may use '-P' for one of SSH-OPTIONs, indicates the port number in
    the ProxyCommand used for the tunneling.
    
EOF
    exit 0
}

declare -A commandmap

commandmap["list"]="list-servers"
commandmap["farms"]="list-farms"
commandmap["set-farm"]="set-farm"
commandmap["env"]="list-environments"
commandmap["set-env"]="set-environment"
commandmap["connect"]="connect-server"
commandmap["help"]="help-and-exit"

while getopts "hvn" opt "$@"; do
    case "$opt" in
        n)
            nocache=1
            ;;
        h)
            help-and-exit
            ;;
        v)
            version-and-exit
            ;;

        f)
            farmid=$OPTARG
            ;;
    esac
done
shift $(($OPTIND - 1))

if [ "$#" -lt 1 ]; then
    error 1 "Try \`-h' for help"
fi

command="$1"
shift

if [ -z "${commandmap[$command]}" ]; then
    error 1 "Try '$PROGRAM_NAME help' for list of commands"
fi
#echo "cmd: ${commandmap[$command]}"

mkdir -p $HOME/.sssh
if [ ! -r "$HOME/.scalr/config.ini" ]; then
    if [ -r "$DEFAULT_SCALR_CONFIG" ]; then
        mkdir $HOME/.scalr
        cp "$DEFAULT_SCALR_CONFIG" $HOME/.scalr/
    else
        error 1 "You need to configure scalr(1) first"
    fi
fi

"${commandmap[$command]}" "$@"