ec2i

#!/bin/bash

set -euo pipefail

CONFIG="$HOME/.ec2i.json"
#if [ -f $HOME/.ec2i-env ]; then
#  . $HOME/.ec2i-env
#fi

OURDIR=$(dirname $0)
OURNAME=$(basename $0)

# ----------------------------------------------------------------
# TODO: usage stdout and exit 0 if -h
usage() {
    echo "Usage: $OURNAME {command}" 1>&2
    cat 1>&2 <<EOF

Given a config file with mappings from epithets/nicknames to AWS EC2 parameters,
allows you to start, stop, ssh, etc. The config file encapsulates instance IDs
and AWS regions, and manages IP address and PEM paths, etc, so you don't have to
keystroke those out all the time.

Requires:

* jq for all processing
* mlr for instance-listing
* github.com/johnkerl/synctool for file transfers

Sample ~/.ec2i.json:

    {
      "pems": [ { "region": "us-east-1", "pem": "/path/to/my.pem" } ],
      "instances": [
        {
          "epithet": "ephemeron",
          "username": "ubuntu",
          "instance_id": "i-0cafefeedcafef00d",
          "region": "us-east-1",
          "description": "This is used for Project 827"
        }
    }

Commands list:

$OURNAME list                 Show epithets and parameters
$OURNAME config               Show the config file with epithet-to-parameter mappings

$OURNAME start    {epithets}  Start the named instance(s)
$OURNAME stop     {epithets}  Stop  the named instance(s)
$OURNAME show     {epithets}  Show config-file parameters for the named instance(s)
$OURNAME ssh      {epithet}   SSH into the named instance (does not auto-start it)

$OURNAME sshup    {epithet}   SSH into the named instance with auto-start and wait for ready

$OURNAME tags     {epithets}  Shows EC2 tags for the named instance(s)
$OURNAME tag      {epithet} {key} {value} Sets an EC2 tag for the named instance
$OURNAME untag    {epithet} {key}         Clears an EC2 tag for the named instance

$OURNAME ip       {epithets}  Shows current IP-address assignment for the named instance(s)
$OURNAME state    {epithets}  Shows one-word EC2 state for the named instance(s)
$OURNAME info     {epithets}  Shows full EC2 JSON info for the named instance(s)

$OURNAME rel-from {epithet} {path}
$OURNAME rel-to   {epithet} {path}
$OURNAME rel-del  {epithet} {path}
$OURNAME pull-pwd {epithet}
$OURNAME push-pwd {epithet}
EOF
    exit 1
}

# ----------------------------------------------------------------
title() {
  text="$@"

  if [ "$(uname)" = "Darwin" ]; then
    starttitle=$'\033]0;'
    endtitle=$'\007'
  else
    starttitle="\e]2;"
    endtitle="\a"
  fi

  echo -e $starttitle $text $endtitle
}

# ================================================================
main() {
    if [ $# -eq 0 ];then
        usage
    fi
    command="$1"
    shift

    # No instance epithet required
    if   [ "$command" = "list" ]; then
        list_instances
        exit 0
    fi
    if   [ "$command" = "config" ]; then
        show_config
        exit 0
    fi

    if [ "$command" = "start" ]; then
        loop_over_epithets $command start_instance "$@"

    elif [ "$command" = "stop" ]; then
        loop_over_epithets $command stop_instance "$@"

    elif [ "$command" = "show" ]; then
        loop_over_epithets $command show_instance "$@"

    elif [ "$command" = "ssh" ]; then
        check_argc_equals $command $# 1 "one epithet from $CONFIG"
        epithet="$1"

        instance_id=$(epithet_to_instance_id $epithet)
        region=$(epithet_to_region $epithet)
        username=$(epithet_to_username $epithet)

        title $epithet
        ssh_instance $username $instance_id $region

    elif [ "$command" = "sshup" ]; then
        check_argc_equals $command $# 1 "one epithet from $CONFIG"
        epithet="$1"

        instance_id=$(epithet_to_instance_id $epithet)
        region=$(epithet_to_region $epithet)
        username=$(epithet_to_username $epithet)

        title $epithet
        start_wait_ssh_instance $username $instance_id $region

    elif [ "$command" = "ip" ]; then
        loop_over_epithets $command get_ip_address "$@"

    elif [ "$command" = "state" ]; then
        loop_over_epithets $command show_instance_state "$@"

    elif [ "$command" = "info" ]; then
        loop_over_epithets $command show_instance_info "$@"

    elif [ "$command" = "tags" ]; then
        loop_over_epithets $command show_instance_tags "$@"

    elif [ "$command" = "tag" ]; then
        check_argc_equals $command $# 3 "epithet, key, and value"
        epithet="$1"
        key="$2"
        value="$3"

        instance_id=$(epithet_to_instance_id $epithet)
        region=$(epithet_to_region $epithet)
        username=$(epithet_to_username $epithet)

        set_instance_tag $instance_id $region $key $value

    elif [ "$command" = "untag" ]; then
        check_argc_equals $command $# 2 "epithet and key"
        epithet="$1"
        key="$2"

        instance_id=$(epithet_to_instance_id $epithet)
        region=$(epithet_to_region $epithet)
        username=$(epithet_to_username $epithet)

        unset_instance_tag $instance_id $region $key

    elif [ "$command" = "rel-from" ]; then
        check_argc_at_least $command $# 2 "epithet and one or more paths"
        epithet=$1; shift
        loop_over_paths "rel-from" "relfromremote" $epithet "$@"

    elif [ "$command" = "rel-to" ]; then
        check_argc_at_least $command $# 2 "epithet and one or more paths"
        epithet=$1; shift
        loop_over_paths $command "reltoremote" $epithet "$@"

    elif [ "$command" = "rel-del" ]; then
        check_argc_at_least $command $# 2 "epithet and one or more paths"
        epithet=$1; shift
        loop_over_paths $command "reldelremote" $epithet "$@"

    elif [ "$command" = "pull-pwd" ]; then
        check_argc_at_least $command $# 1 "epithet and no paths"
        epithet=$1; shift
        loop_over_paths $command "pullpwd" $epithet "$@"

    elif [ "$command" = "push-pwd" ]; then
        check_argc_at_least $command $# 1 "epithet and no paths"
        epithet=$1; shift
        loop_over_paths $command "pushpwd" $epithet "$@"

    else
        usage
    fi
}

# ----------------------------------------------------------------
check_argc_equals() {
    command="$1"
    actual_argc="$2"
    expected_argc="$3"
    needed="$4"

    if [ $actual_argc -ne $expected_argc ]; then
        echo "$OURNAME $command: need $needed." 1>&2
        exit 1
    fi
}

check_argc_at_least() {
    command="$1"
    actual_argc="$2"
    expected_argc="$3"
    needed="$4"

    if [ $actual_argc -lt $expected_argc ]; then
        echo "$OURNAME $command: need $needed." 1>&2
        exit 1
    fi
}

loop_over_epithets() {
    user_facing_command="$1"
    function_name="$2"
    shift; shift
    epithets="$@"

    if [ $# -lt 1 ]; then
        echo "$OURNAME $user_facing_command: need at least one instance epithet from $CONFIG." 1>&2
        exit 1
    fi
    i=0
    for epithet; do
        # Put a blank line between outputs for multiple instances
        if [ $# -gt 1 -a $i -gt 0 ]; then
            echo
        fi
        i=$[i+1]

        instance_id=$(epithet_to_instance_id $epithet)
        region=$(epithet_to_region $epithet)
        username=$(epithet_to_username $epithet)
        $function_name $instance_id $region
    done
}

loop_over_paths() {
    user_facing_command="$1"
    synctool_function_name="$2"
    epithet="$3"
    shift; shift; shift

    instance_id=$(epithet_to_instance_id $epithet)
    region=$(epithet_to_region $epithet)
    username=$(epithet_to_username $epithet)

    synctool_instance $username $instance_id $region $synctool_function_name "$@"
}

# ================================================================
list_instances() {
    regions=$(jq -r .instances[].region $CONFIG | sort -u)
    for region in $regions; do
        AWS_REGION=$region AWS_DEFAULT_REGION=$region aws ec2 describe-instances
    done | mlr --j2p put -q '

        for (reservation in $Reservations) {
            for (instance in reservation["Instances"]) {
                KeyName = instance["KeyName"];
                if (haskey(instance, "KeyName") && index(KeyName, "kerl") >= 0) {
                    tags = instance["Tags"];

                    # [
                    #   {
                    #     "Key": "Name",
                    #     "Value": "kerl-tiny"
                    #   },
                    #   {
                    #     "Key": "Epithet",
                    #     "Value": "dwergaz"
                    #   }
                    # ]

                    tags = sort(tags, func(a,b) { return a["Key"] <=> b["Key"]});
                    # tags = select(tags, func(e) { return e["Key"] != "Epithet" && e["Key"] != "Name"});
                    # tag_values = apply(sort(tags), func(e) { return e["Value"] });

                    out = {};
                    # out["KeyName"] = KeyName;

                    epithet = "???";
                    name = "???";
                    os = "";
                    arch = "";
                    cpu = "";
                    ram = "";
                    pph = "";

                    for (tag in tags) {
                        # print "TAG:", tag;
                        if (tag["Key"] == "Epithet") {
                            epithet = tag["Value"];
                        } elif (tag["Key"] == "Name") {
                            name = tag["Value"];

                        } elif (tag["Key"] == "OS") {
                            os = tag["Value"];
                        } elif (tag["Key"] == "Arch") {
                            arch = tag["Value"];
                        } elif (tag["Key"] == "CPU") {
                            cpu = int(tag["Value"]);
                        } elif (tag["Key"] == "RAM") {
                            ram = int(tag["Value"]);
                        } elif (tag["Key"] == "PPH") {
                            pph = int(tag["Value"]);
                        }
                    }

                    out["State"]   = instance["State"]["Name"];
                    out["Epithet"] = epithet;
                    out["Name"]    = name;
                    out["OS"]      = os;
                    out["Arch"]    = arch;
                    out["CPU"]     = cpu;
                    out["RAM"]     = ram;
                    out["PPH"]     = pph;

                    out["InstanceId"]   = instance["InstanceId"];
                    out["InstanceType"] = instance["InstanceType"];

                    emit out;
                }
            }
        }

    ' then sort -n State,PPH,Epithet
}

show_config() {
    jq '.pems[]' $CONFIG | mlr --j2p cat
    echo
    jq '.instances[]' $CONFIG | mlr --j2p cat
}

# ----------------------------------------------------------------
epithet_to_instance_id() {
    epithet="$1"
    instance_id=$(jq -r '.instances[] | select(.epithet=="'$epithet'") | .instance_id' $CONFIG)
    if [ "$instance_id" = "" ]; then
        echo "$OURNAME: could not find instance_id for epithet $epithet within config $CONFIG" 1>&2
        exit 1
    fi
    echo $instance_id
}

epithet_to_region() {
    epithet="$1"
    region=$(jq -r '.instances[] | select(.epithet=="'$epithet'") | .region' $CONFIG)
    if [ "$region" = "" ]; then
        echo "$OURNAME: could not find region for epithet $epithet within config $CONFIG" 1>&2
        exit 1
    fi
    echo $region
}

epithet_to_username() {
    epithet="$1"
    username=$(jq -r '.instances[] | select(.epithet=="'$epithet'") | .username' $CONFIG)
    if [ "$username" = "" ]; then
        echo "$OURNAME: could not find username for epithet $epithet within config $CONFIG" 1>&2
        exit 1
    fi
    echo $username
}

# ----------------------------------------------------------------
show_instance() {
    instance_id="$1"
    region="$2"
    AWS_REGION=$region AWS_DEFAULT_REGION=$region \
        aws ec2 describe-instances --query "Reservations[0].Instances[0]" --instance-ids $instance_id \
            | mlr --ijson --oxtab \
              cut -o -f KeyName,InstanceType,InstanceId,Architecture,LaunchTime,Placement.AvailabilityZone,State
}

start_instance() {
    instance_id="$1"
    region="$2"
    AWS_REGION=$region AWS_DEFAULT_REGION=$region aws ec2 start-instances --instance-ids $instance_id
}

stop_instance() {
    instance_id="$1"
    region="$2"
    AWS_REGION=$region AWS_DEFAULT_REGION=$region aws ec2 stop-instances --instance-ids $instance_id
}

ssh_with_pem() {
    username_and_ip_address="$1"
    pem_file="$2"
    echo
    echo ssh -i $pem_file $username_and_ip_address

    ok="true"
    ssh -i $pem_file $username_and_ip_address || ok="false"
    echo "ssh ok $ok"

    if [ $ok != "true" ]; then
      # Error message already printed out
      exit 1
    fi
}

get_ip_address() {
    instance_id="$1"
    region="$2"
    AWS_REGION=$region AWS_DEFAULT_REGION=$region aws ec2 describe-instances --instance-ids $instance_id \
        | jq -r ".Reservations[0].Instances[0].PublicIpAddress"
}

get_pem_file() {
    region="$1"
    pem=$(jq -r '.pems[] | select(.region=="'$region'") | .pem' $CONFIG)
    if [ "$pem" = "null" ]; then
        echo "$OURNAME: could not find PEM for region $region within config $CONFIG"
        exit 1
    fi
    echo $pem
}

ssh_instance() {
    username="$1"
    instance_id="$2"
    region="$3"

    echo
    echo "username    $username"
    echo "instance_id $instance_id"
    echo "region      $region"

    ip_address=$(get_ip_address $instance_id $region)
    echo "ip_address  $ip_address"

    pem_file=$(get_pem_file $region)
    echo "pem_file    $pem_file"
    echo

    if [ "$ip_address" = "null" ]; then
        echo "Could not find IP address for $instance_id in $region. Is it running?"
    else
        ssh_with_pem $username@$ip_address $pem_file
    fi
}

start_wait_ssh_instance() {
    username="$1"
    instance_id="$2"
    region="$3"

    state=$(show_instance_state $instance_id $region)

    echo
    echo "username    $username"
    echo "instance_id $instance_id"
    echo "region      $region"
    echo "state       $state"

    pem_file=$(get_pem_file $region)
    echo "pem_file    $pem_file"
    echo

    if [ "$state" != "running" ]; then

        # Wait for EC2 spin-up
        start_instance $instance_id $region
        n=1
        while true; do
            sleep 1
            n=$[n+1]
            if [ $n -gt 30 ]; then
                echo "$0: giving up on start of $instance_id in $region"
                exit 1
            fi
            state=$(show_instance_state $instance_id $region)
            echo "state       $state"
            if [ "$state" = "running" ]; then
                break
            fi
        done
    fi

    # Wait for sshd etc.
    n=1
    while true; do
        sleep 1
        n=$[n+1]
        if [ $n -gt 20 ]; then
            echo "$0: giving up on ssh to $instance_id in $region"
            exit 1
        fi

        ip_address=$(get_ip_address $instance_id $region)
        if [ "$ip_address" = "null" ]; then
            echo "Could not find IP address for $instance_id in $region. Is it running?"
            continue
        fi

        t1=$(date +%s)
        echo "timestamp   $t1"

        echo ssh -i $pem_file $username@$ip_address
        ok="true"

        ssh -i $pem_file $username@$ip_address || ok="false"

        t2=$(date +%s)
        dt=$[t2 - t1]

        echo "ssh ok      $ok"
        echo "timestamp   $t2"
        echo "ssh seconds $dt"

        if [ $ok == "true" ]; then
          break
        fi
        if [ $dt -gt 100 ]; then
          # Example: voluntary "~." hangup after long-running session.
          # Don't retry. Our retry loop here is intended solely for
          # adaptive timing in instance-startup scenarios.
          break
        fi
        echo "Awaiting ssh ..."
    done
}

synctool_instance() {
    username="$1"
    instance_id="$2"
    region="$3"
    op="$4"
    shift; shift; shift; shift

    echo
    echo "instance_id $instance_id"
    echo "region      $region"

    ip_address=$(get_ip_address $instance_id $region)
    echo "ip_address  $ip_address"

    pem_file=$(get_pem_file $region)
    echo "pem_file    $pem_file"
    echo

    if [ "$ip_address" = "null" ]; then
        echo "Could not find IP address for $instance_id in $region. Is it running?" 1>&2
    else
        synctool -i $pem_file -r $username@$ip_address $op "$@"
    fi
}

# ----------------------------------------------------------------
show_instance_state() {
    instance_id="$1"
    region="$2"
    AWS_REGION=$region AWS_DEFAULT_REGION=$region aws ec2 describe-instances --instance-ids $instance_id \
        | jq -r ".Reservations[0].Instances[0].State.Name"
}

show_instance_info() {
    instance_id="$1"
    region="$2"
    AWS_REGION=$region AWS_DEFAULT_REGION=$region aws ec2 describe-instances --instance-ids $instance_id | jq .
}

show_instance_tags() {
    instance_id="$1"
    region="$2"
    AWS_REGION=$region AWS_DEFAULT_REGION=$region aws ec2 describe-tags --filters Name=resource-id,Values=$instance_id \
        | jq .Tags | mlr --j2p cat
}

set_instance_tag() {
    instance_id="$1"
    region="$2"
    name="$3"
    value="$4"
    AWS_REGION=$region AWS_DEFAULT_REGION=$region aws ec2 create-tags --resources $instance_id --tags Key="$name",Value="$value"
}

unset_instance_tag() {
    instance_id="$1"
    region="$2"
    name="$3"
    AWS_REGION=$region AWS_DEFAULT_REGION=$region aws ec2 delete-tags --resources $instance_id --tags Key="$name"
}

# ================================================================
main "$@"