#! /usr/bin/bash

# Start AWS instance easily, all-or-nothing, with Resalloc API
# Copyright (C) 2021 Pavel Raiskup <praiskup@redhat.com>

# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

set -e

scriptdir=$(dirname "$(readlink -f "$0")")

show_help()
{
cat <<EOHELP >&2
Usage: $0 [OPTIONS]

Script is aimed to help sysadmin.

Instead of --name, you can use \$RESALLOC_NAME env var.

Options:
  --name NAME           Name of the started VM instance
  --ami IMAGE           Start the VM from this AMI.
  --aws-profile         Name of the aws profile, as defined in ~/.aws/config.
  --ssh-key-name        Name of the uploaded SSH key uploaded to AWS.
  --security-group-id SGID
                        ID of the AWS security group to place this in (AWS
                        firewall), e.g. 'sg-0c3efdb676ced6d4f'.
  --instance-type INSTANCE_TYPE_ID
                        The instance type of the instance to use, "flavor".
                        Defines the parameters of the machine (CPU, ...),
                        e.g. 'i3.large'.
  --possible-subnet SUBNET_ID
                        Allow using the given SUBNET_ID, for example
                        subnet-08d0b3168a353e3ee.  Required option.  This can
                        option can be used multiple times; if so - the allocated
                        instance uses just one of the given subnets, picked by
                        random.
  --tag TAGNAME=VALUE
                        Tag every resource started in EC2 with TAGNAME=VALUE.
                        Can be used multiple times.
  --initial-preparation Start the vm and perform an initial set of actions
                        (using playbook) on the fresh Fedora Cloud image (given
                        by AMI_ID) so we can later generate a new AMI from that
                        VM.  See also --create-snapshot-image.
  --create-snapshot-image
                        Create a new AMI image from the started VM.  Note that
                        when --additional-volume-size is used - such volume is
                        inherited and then automatically created upon
                        instatiating the image.
  --spot-price PRICE    Start a spot instance, instead of (default) on-demand,
                        with the given maximal price per hour.
  --playbook PLAYBOOK   Use this playbook, instead of the pre-configured one.
  --root-volume-size GB
                        Allocate root volume of given size.
  --additional-volume-size GB
                        Allocate additional volume of given size.
  --private-ip          This VM starts on private IP, work over private IP.
  --no-print-ip         Don't print the detect IP on stdout (once finished).
                        This is the only output on stdout, and resalloc server
                        requires it.
  --debug               Print debugging info.
EOHELP

test -z "$1" || exit "$1"
}

opt_debug=false

info() { echo >&2 " * $*" ; }
debug() { if $opt_debug; then echo >&2 " - $*" ; fi }
error() { echo -e >&2 "!!!\n!!! $*\n!!!\n"; }
fatal_args () { error "$*"; show_help 1 ; }

run_cmd()
{
    debug "run_cmd: $*"
    "$@"
}

# handle no arguments
test "${#@}" -eq 0 && show_help 1

long_opts="name:,ami:,initial-preparation,create-snapshot-image,help,spot-price:,\
playbook:,private-ip,debug,root-volume-size:,additional-volume-size:,no-print-ip,aws-profile:,\
ssh-key-name:,security-group-id:,possible-subnet:,instance-type:,tag:"
ARGS=$( getopt -o "h" -l "$long_opts" -n "getopt" -- "$@") || show_help 1

mandatory_options="
    --ami
    --aws-profile
    --instance-type
    --ssh-key-name
    --security-group-id
"

# is that necessary -> should preserve whitespaces in option arguments
# see: http://linuxwell.com/2011/07/14/getopt-in-bash/
eval set -- "$ARGS"

opt_initial_preparation=false
opt_root_volume_size=
opt_additional_volume_size=
opt_create_snapshot_image=false
opt_ami=
opt_aws_profile=
opt_name=$RESALLOC_NAME
opt_spot_price=
opt_playbook=
opt_private_ip=false
opt_print_ip=true
opt_ssh_key_name=
opt_security_group_id=
opt_possible_subnet=()
opt_tag=()

instance_id=

option_variable()
{
    opt=$1
    opt=${1##--}
    opt=${opt##-}
    opt=${opt//-/_}
    # drop the '--no' prefix
    if test -n "$2"; then
        opt=${opt//no_/}
    fi
    option_variable_result=opt_$opt
}

while true; do
    # now the name is in $1 and argument in $2
    case $1 in
    -h|--help)
        show_help 0
        ;;

    --name|--ami|--spot-price|--playbook|--root-volume-size|--additional-volume-size|\
    --aws-profile|--ssh-key-name|--security-group-id|--instance-type)
        option_variable "$1"
        eval "$option_variable_result=\$2"
        shift 2
        ;;

    --initial-preparation|--create-snapshot-image|--private-ip|--debug)
        option_variable "$1"
        eval "$option_variable_result=:"
        shift
        ;;

    --no-print-ip)
        option_variable "$1"
        eval "$option_variable_result=false"
        shift
        ;;

    # Array options
    --possible-subnet|--tag)
        option_variable "$1"
        eval "opt_value=( \"\${$option_variable_result[@]}\" )"
        opt_value+=( "$2" )
        eval "$option_variable_result"'=( "${opt_value[@]}" )'
        shift 2
        ;;

    --) shift; break ;;  # end
    *) echo "programmer mistake ($1)" >&2; exit 1 ;;
    esac
done

test -z "$opt_name" && fatal_args '$RESALLOC_NAME or --name required'

test -z "${opt_possible_subnet[*]}" && \
    fatal_args "at least one --possible-subnet required"

for mandatory_option in $mandatory_options
do
    option_variable "$mandatory_option"
    eval "test -z \$$option_variable_result" && \
        fatal_args "$mandatory_option is required"
done

aws=(
    aws --profile "$opt_aws_profile"
        --output text
)

debug "AWS base command '${aws[*]}'"

info "Instance Name: $opt_name"

tmpdir=/tmp
workdir=$(mktemp -d "$tmpdir"/resalloc-aws-start-attempt-XXXXXXXX)
logfile=$workdir/playbook.log
cleanup_vm=true
cleanup()
{
    if $cleanup_vm && test -n "$instance_id"; then
        # cleanup the VM we failed to start, but keep logs
        info "removing $instance_id"
        run_cmd "${aws[@]}" ec2 terminate-instances --instance-ids "$instance_id"
    else
        # Remove the logs in case of *no* problem happened ...
        case $workdir in
            /tmp/*) rm -rf "$workdir" ;;
        esac
    fi
}
trap cleanup EXIT

tagspec()
{
    tagsep=
    tagspec_result=
    for tag in "$@"; do
        old_IFS=$IFS
        IFS='='
        set -- $tag
        tagspec_result+="${tagsep}{Key=$1,Value=$2}"
        tagsep=,
        IFS=$old_IFS
    done
}

tagspec "Name=$opt_name" "${opt_tag[@]}"

start_opts=(
    --key-name "$opt_ssh_key_name"
    --security-group-ids "$opt_security_group_id"
    --tag-specifications "ResourceType=instance,Tags=[$tagspec_result]"
                         "ResourceType=volume,Tags=[$tagspec_result]"
    --count 1
    --query 'Instances[0].InstanceId'
)

subnets=()
start_opts+=(
    --image-id "${opt_ami}"
    --instance-type "${opt_instance_type}"
)
subnets=( "${opt_possible_subnet[@]}" )

# We pick a random subnet from the given options, this brings a lower chance of
# subsequent script failure in case of some temporary datacenter failure.
subnet_count=${#subnets[@]}
pick_subnet=$(( RANDOM % subnet_count ))
eval 'subnet_id=${subnets['"$pick_subnet"']}'
start_opts+=( --subnet-id "$subnet_id" )

info "Subnet ID: $subnet_id (from ${subnets[*]})"

if test -n "$opt_spot_price"; then
    market_options="MarketType=spot,SpotOptions={MaxPrice=$opt_spot_price,\
InstanceInterruptionBehavior=terminate,SpotInstanceType=one-time}"
    start_opts+=( --instance-market-options "$market_options" )
fi

# Use gp3 volume for the root filesystem
device_mappings='{"DeviceName":"/dev/sda1","Ebs":{"VolumeType":"gp3"'

if test -n "$opt_root_volume_size"; then
    device_mappings+=',"VolumeSize":'
    device_mappings+=$opt_root_volume_size
fi

device_mappings+='}}'

if test -n "$opt_additional_volume_size"; then
    device_mappings+=',{"DeviceName":"/dev/sdd","Ebs":{"VolumeSize":'
    device_mappings+=$opt_additional_volume_size
    device_mappings+=',"VolumeType":"gp3","DeleteOnTermination":true}}'
fi

start_opts+=( --block-device-mappings "[$device_mappings]" )

aws_cmd=( "${aws[@]}" ec2 run-instances "${start_opts[@]}" )
instance_id=$( run_cmd "${aws_cmd[@]}" )
info "Instance ID: $instance_id"
test -n "$instance_id"

ip_type=Public
$opt_private_ip && ip_type=Private

info "Waiting till the IPv4 is assigned"
for _ in $(seq 100); do
ip_address=$(
    run_cmd "${aws[@]}" ec2 describe-instances \
    --query "Reservations[*].Instances[*].${ip_type}IpAddress" \
    --instance-ids "$instance_id"
)
test -n "$ip_address" && break
sleep 2
done

case $ip_address in
    *.*.*.*) ;;
    *) info "bad ip address: '$ip_address'" ; false ;;
esac

info "Instance IP: $ip_address"

info "waiting for ssh ..."

run_cmd "$scriptdir/resalloc-wait-for-ssh" \
    --check-cloud-user --timeout 300 "$ip_address" \
    --log debug

if test -n "$opt_playbook"; then
    playbook_opts=()
    $opt_initial_preparation && playbook_opts+=( "-e" "prepare_base_image=1" )
    run_cmd ansible-playbook -i "$ip_address," "$opt_playbook" "${playbook_opts[@]}" >&2 </dev/null
fi

if $opt_create_snapshot_image; then
    new_image_name=$opt_name-$(date +"%Y%m%d_%H%M%S")
    tagspec "Name=$new_image_name" "${opt_tag[@]}"
    image_id_cmd=(
        "${aws[@]}" ec2 create-image
        --instance-id "$instance_id"
        --name "$new_image_name"
        --output text
        --tag-specifications "ResourceType=image,Tags=[$tagspec_result]"
    )

    info "Preparing image $new_image_name"
    image_id=$( run_cmd "${image_id_cmd[@]}" )
    test -n "$image_id"
    info "Image ID: $image_id"
    "${aws[@]}" ec2 wait image-available --image-ids "$image_id"

else
    cleanup_vm=false
fi

$opt_print_ip && echo "$ip_address"
