run-and-connect

#!/bin/bash
#
# Script for launching the specified Amazon Machine Image (AMI) and connecting
# to it via SSH. Automatically terminates upon logout from the EC2 instance.
#
# See https://github.com/susestudio/susestudio-ec2/blob/master/README.md for
# details.
#
# Author: James Tan <jatan@suse.com>

region=eu-west-1

function usage() {
  echo "Usage: run-and-connect [--region REGION] AMI_ID"
  echo
  echo "Launches and connects to the specified AMI in the specified region."
  echo
  echo "Script creates a one-time SSH keypair, waits for the AMI to be available,"
  echo "launches the instance, waits for it to boot, then SSHs into it. It"
  echo "automatically terminates the EC2 instanceand removes the created SSH"
  echo "keypair on logout."
  echo
  echo "Options:"
  echo "  --region REGION      The region to upload and register in [us-east-1, us-west-1,"
  echo "                       eu-west-1, ap-southeast-1, ap-northeast-1]. Default is '$region'."
  echo "  --self_test          Run self test on AMI instance and exit."
  echo "  --help               Display this help and exit."
}

function create_ssh_keypair() {
  local random=`mktemp -u | cut -d'.' -f2`
  ssh_key_name="suse-studio.$random"

  log "Creating SSH keypair $ssh_key_name in $region..."
  ec2-add-keypair $ssh_key_name --region $region > $ssh_key_name
  exit_if_failed

  chmod 600 $ssh_key_name
}

function start_instance() {
  log "Starting new instance ($root_device, $arch)..."
  if [ "$root_device" = "instance-store" ]; then
    instance_type=m1.small
    [ "$arch" = x86_64 ] && instance_type=m1.large
  else
    instance_type=t1.micro
  fi
  output=`ec2-run-instances $ami --region "$region" -k "$ssh_key_name" -g "SUSE_Studio" -t $instance_type 2>&1`
  exit_if_failed
  instance=`echo $output | cut -d' ' -f6`
  if [ ! "$instance" ]; then
    log "Failed to start instance:"
    echo $output
    exit 1
  fi
  log "Started $instance"
}

function wait_for_ami() {
  log "Waiting for AMI..." -n
  for i in {1..100}; do
    output=`ec2-describe-images --region "$region" "$ami" 2>&1`
    exit_if_failed
    state=`echo $output | cut -d' ' -f5`
    arch=`echo $output | cut -d' ' -f7`
    root_device=`echo $output | cut -d' ' -f10`
    [ "$state" ] && [ ! "$state" = "pending" ] && break
    sleep 3
    echo -n "."
  done
  echo ""
  if [ ! "$state" ] || [ ! "$state" = "available" ]; then
    log "Failed to get available AMI"
    exit 1
  fi
}

function wait_for_hostname() {
  log "Waiting for hostname..." -n
  for i in {1..10}; do
    output=`ec2-describe-instances --region "$region" "$instance" 2>&1`
    exit_if_failed
    hostname=`echo $output | cut -d' ' -f8`

    if [ "$hostname" ] && [ ! "$hostname" = "pending" ]; then
      break
    fi
    sleep 3
    echo -n "."
  done
  echo ""
  if [ ! "$hostname" ] || [ "$hostname" = "pending" ] ; then
    log "Failed to get hostname (timed out)"
    exit 1
  fi
  log "Hostname is $hostname"
}

function wait_for_ssh() {
  local ssh_err=.ssh_error
  log "Waiting for instance boot and SSH..." -n
  output=""
  for i in {1..50}; do
    output=`ssh -i $ssh_key_name -o "StrictHostKeyChecking no" -o "UserKnownHostsFile /dev/null" -o "BatchMode yes" -o "ConnectTimeout 10" root@$hostname "echo 'hi'" 2>$ssh_err`
    [ "$output" = "hi" ] && break
    sleep 3
    echo -n "."
  done
  echo ""

  # Check if we have timed out
  if [ ! "$output" = "hi" ]; then
    log "Failed to connect via SSH (timed out)"
    [ -f $ssh_err ] && cat $ssh_err
    clean_up
    exit 1
  fi
}

function exit_if_failed() {
  local code=$?
  local newline=$1
  if [ $code -ne 0 ]; then
    [ "$newline" ] && echo ""
    die $code "$output"
  fi
}

function die() {
  local code=$1
  local output=$2
  log "Failed (exit code $code)"
  [ "$output" ] && echo "$output"
  clean_up
  exit 10
}

function log() {
  local msg=$1
  local opts=$2
  local time=`date +%H:%M:%S`
  echo $opts "$time $msg"
}

function clean_up() {
  log "Cleaning up..."
  if [ "$ssh_key_name" ]; then
    log "Removing SSH keypair $ssh_key_name..."
    output=`ec2-delete-keypair $ssh_key_name --region "$region" 2>&1`
    if [ $? -ne 0 ]; then
      log "Failed to remove SSH keypair"
      echo $output
    fi
    rm -f $ssh_key_name
  fi
  if [ "$instance" ]; then
    log "Terminating instance $instance..."
    output=`ec2-terminate-instances --region "$region" $instance 2>&1`
    if [ $? -ne 0 ]; then
      log "Termination failed."
    fi
  fi
  log "Bye"
}

function ssh_exec() {
  local cmd=$1
  log "ssh_exec: $cmd"
  ssh -q -i $ssh_key_name -o "StrictHostKeyChecking no" -o "UserKnownHostsFile /dev/null" root@$hostname "$cmd"
  local code=$?
  if [ $code -ne 0 ]; then
    log "ssh_exec failed with exit code $code"
    exit $code
  fi
}

# Parse and set options
run_self_test=false
while [ $# -gt 0 ]; do
  case "$1" in
    --region)    shift; region="$1";;
    --self_test) run_self_test=true;;
    -h|--help)   usage; exit 1;;
    *) break;;
  esac
  shift
done

if [ $# -ne 1 ]; then
  usage
  exit 1
fi

ami=$1

create_ssh_keypair
wait_for_ami
start_instance
wait_for_hostname
wait_for_ssh

if [ "$run_self_test" = true ]; then
  ssh_exec "zypper -n refresh"
else
  ssh -i $ssh_key_name -o "StrictHostKeyChecking no" -o "UserKnownHostsFile /dev/null" root@$hostname
fi

clean_up
exit 0