~smoser/ubuntu-on-ec2/ec2-publishing-scripts.lp927823

#!/bin/bash
# vi: ts=4 noexpandtab

PREFIX=""
TEMP_D=""
SSH_USER="ubuntu"
IID=""
error() { echo ${PREFIX:+"${PREFIX}"} "$@" 1>&2; }
errorp() { printf "$@" 1>&2; }
fail() { 
	[ $# -eq 0 ] || error "$@";
	if [ -n "${IID}" ]; then
		local regargs=()
		[ -n "${REGION}" ] && regargs=( --region "${REGION}" )
		error "terminating ${IID}"
		error "===== console log begin ====="
		xc2 get-console-output "${regargs[@]}" "${IID}" 1>&2
		error "===== console log end   ====="
		error "  "
		error "===== instance properties begin ===="
		xc2 info-instances "${regargs[@]}" 1>&2
		error "===== instance properties end ===="
		error " "
		error "===== elastic ip assignments begin ===="
		xc2 describe-addresses "${regargs[@]}" 1>&2
		error "===== elastic ip assignments end ===="
		xc2 terminate-instances "${regargs[@]}" "${IID}"
	fi
	exit 1;
}
failp() { [ $# -eq 0 ] || errorp "$@"; fail; }
debug() {
	local level=${1}
	shift;
	[ "${level}" -gt "${VERBOSITY}" ] && return
	error "$(date)" "${@}"
}
cleanup() {
	[ -z "${TEMP_D}" ] || rm -Rf "${TEMP_D}"
}
sigint() { error "caught user interupt"; fail; }
sshconnect() {
	local out
	out=$(ssh -o ConnectTimeout=5 -o BatchMode=yes "$@" "/bin/sh -c 'echo true'")
	[ -n "$out" -a "$out" != "true" ] && error "unexpected output: $out"
	[ "${out}" = "true" ]
}

Usage() {
	cat <<EOF
Usage: ${0##*/} [ options ] ami-id region 

   Run instance of 'ami-id' in 'region'.  Wait for it to come up
   and verify ssh connection before returning

   Will exit failure if instance seems not to have come up.
   It should not exit failure without killing instance.

   options:
   -h|--help                  print this help
      --instance-type x       start instance of instance-type x
      --key           x       start instance with --key x
      --ssh-pubkey    f       connect with 'ssh -i f'
      --associate-ip  i       associate elastic IP i with instance
      --known-hosts   f       write ssh fingerprint to ssh known hosts
                              file f, removing old entries for this host
      --prefix        p       all debug/status messages prefixed with 'p'
EOF
}

bad_Usage() { Usage 1>&2; [ $# -eq 0 ] || error "$@"; exit 1; }

di_field() {
	local field=$1 f=""
	local fields=( itype iid ami host host_int state )
	for((i=0;i<${#fields[@]};i++)); do
		[ "${fields[$i]}" = "${field}" ] && { f=$(($i+1)) ; break; }
	done
	[ -n "$f" ] || return 1
	local cmd='$1 == "INSTANCE" { print $'$f' }'
	_RET=$(awk '-F\t' "$cmd")
}

write_ud_script() {
	local out_sshd_pub=$1 out_script=$2 user_pubkey=$3
	local ssh_user=${SSH_USER:-ubuntu} 
	ssh-keygen -q -C "root@cloudhost" \
		-f "${TEMP_D}/ssh_host_rsa_key" -t rsa -N "" ||
		{ error "failed to gen ssh key"; return 1; }
	{
	cat <<EOF
#!/bin/sh

echo ============ \$(date) =============
echo ############ populating keys ##############
(
umask 066
cat >/etc/ssh/ssh_host_rsa_key <<END
$(cat "${TEMP_D}/ssh_host_rsa_key")
END
)

cat >/etc/ssh/ssh_host_rsa_key.pub <<END
$(cat "${TEMP_D}/ssh_host_rsa_key.pub")
END
ssh-keygen -l -f /etc/ssh/ssh_host_rsa_key.pub
EOF
if [ -n "${user_pubkey}" ]; then
cat <<END_CAT
sudo -Hu "${ssh_user}" sh <<"END_SH"
d=\$HOME/.ssh; mkdir -p "\$d" 2>/dev/null && chmod 700 "\$d" &&
cat >> "\${d}/authorized_keys" <<"END_KEY"
$(cat "${user_pubkey}")
END_KEY
END_SH
END_CAT
fi
true
	} > "${out_script}"
	[ $? -eq 0 ] || { error "failed to write to ${out_script}"; return 1; }

	cat "${TEMP_D}/ssh_host_rsa_key.pub" > "${out_sshd_pub}" ||
		{ error "failed to write to ${out_sshd_pub}"; return 1; }
}

short_opts="hv"
long_opts="associate-ip:,help,instance-type:,key:,known-hosts:,prefix:,region:,ssh-pubkey:,verbose"
getopt_out=$(getopt --name "${0##*/}" \
	--options "${short_opts}" --long "${long_opts}" -- "$@") &&
	eval set -- "${getopt_out}" ||
	bad_Usage

ami=""
host=""
sleeptime=10
region=""
instance_type="m1.small"
keyname=""
ssh_wait_max=300
ssh_pubkey=""
associate_ip=""
known_hosts="${HOME}/.ssh/known_hosts"
pt=( )
PREFIX=""

while [ $# -ne 0 ]; do
	cur=${1}; next=${2};
	case "$cur" in
		-h|--help) Usage; exit 0;;
		-v|--verbose)
			VERBOSITY=$((${VERBOSITY}+1));;
		   --instance-type)
			pt[${#pt[@]}]="$cur"; pt[${#pt[@]}]="$next"
			instance_type=${next}; shift;;
           --key)
			pt[${#pt[@]}]="$cur"; pt[${#pt[@]}]="$next"
			keyname=${next}; shift;;
		   --ssh-pubkey)
			ssh_pubkey="${next}"; shift;;
		   --associate-ip)
		    associate_ip="${next}"; shift;;
		   --known-hosts)
			known_hosts="${next}"; shift;;
		   --prefix)
			PREFIX="${next}"; shift;;
		--) shift; break;;
	esac
	shift;
done

[ $# -eq 1 -o $# -eq 2 ] || bad_Usage "must provide ami"
ami=${1}
region=${2}

[ -n "${region}" ] &&
    REG_ARGS=( "--region" "${region}" )

START_TIME=$SECONDS
TEMP_D=$(mktemp -d ${TMPDIR:-/tmp}/${0##*/}.XXXXXX) ||
	fail "failed to make tmp dir"
trap cleanup EXIT
trap sigint SIGINT

out_inst_run="${TEMP_D}/instance-run.out"
err_inst_run="${TEMP_D}/instance-run.err"
out_inst_desc="${TEMP_D}/instance-describe.out"
out_keyscan="${TEMP_D}/keyscan.out"
out_userdata="${TEMP_D}/user-data.script"
out_sshd_pub="${TEMP_D}/sshd_pub.txt"
out_associate_ip="${TEMP_D}/out_associate_ip.txt"
out_clean_knownhosts="${TEMP_D}/out_clean_knownhosts.txt"
connect_err="${TEMP_D}/ssh_connect_err.txt"

write_ud_script "${out_sshd_pub}" "${out_userdata}" \
	${ssh_pubkey:+"${ssh_pubkey}"} ||
	fail "failed to write user data script"

debug 1 "xc2 run-instances ${REG_ARGS[@]} ${pt[@]} ${ami}"
insuff_retries=24
for((i=0;i<${insuff_retries};i++)); do
	xc2 run-instances "${REG_ARGS[@]}" "${pt[@]}" \
		"--user-data-file=${out_userdata}" "$ami" \
			> "${out_inst_run}" 2>"${err_inst_run}" &&
		break
	if grep -q "Server.InsufficientInstanceCapacity" "${err_inst_run}"; then
		debug 1 "Server.InsufficientInstanceCapacity, sleeping 5m and trying again [${i}/${insuff_retries}]"
		sleep 5m
	else
		cat "${err_inst_run}" 1>&2
		fail "failed to run instance"
	fi
done
start=${SECONDS}

di_field iid < "${out_inst_run}" && iid=${_RET} && [ -n "${iid}" ] ||
	fail "reading iid from run-instances output failed"

# IID and REGION read by fail code
IID="${iid}"
REGION=${region}

debug 1 "instance id is ${iid}"

if [ -n "${associate_ip}" ]; then
	n=0
	while ! xc2 associate-address "${REG_ARGS[@]}" \
		-i "${iid}" "${associate_ip}" \
		> "${out_associate_ip}" 2>&1 &&
		[ $n -lt 5 ]; do
		debug 2 "associate-ip -i ${iid} ${associate_ip} failed, trying again"
		sleep 3
		n=$((n+1))
	done
	[ ${n} -eq 5 ] && {
		cat "${out_associate_ip}" 1>&2;
		fail "failed to associate ip ${associate_ip}";
	}
	debug 1 "associated ip ${associate_ip} to ${iid}"
fi

while sleep ${sleeptime}; do
	xc2 info-instances "${REG_ARGS[@]}" "${iid}" > "${out_inst_desc}" &&
		di_field state < "${out_inst_desc}" ||
		fail "info-instances ${iid} failed"
	state=${_RET}
	case ${state} in
		pending) :;;
		terminated) fail "${iid} went to terminated";;
		running) break;;
		*) fail "${iid} in unknown state: ${_RET}";;
	esac
	debug 2 "${iid} in ${state} ($((${SECONDS}-${start})) seconds)"
done

di_field host < "${out_inst_desc}" && [ -n "${_RET}" ] &&
	host=${_RET} || fail "host running, but no hostname!"

if [ -n "$associate_ip" ]; then
	# After associating an IP address, there is a period of time when
	# describe-images will still show the old (now unusable) hostname
	#
	# in probably 90% of cases, by the time we got here, the final
	# describe-images gave the updated host, but sometimes it would still
	# have the old host.  The result was that we would then attempt to
	# address the instance by the old/invalid hostname and the ssh-keyscan
	# would fail.
	#
	# so, here we just make hostname == associated_ip, as that should
	# work in all cases below.
	host="${associate_ip}"
fi

debug 1 "instance as $iid/${host} running after $((${SECONDS}-${start}))s"

# here you could/should do either
# - scrape console for keys
# - verify keys are as expected due to stuffing in user-data or cloud-config
waited=0
ssh_start=${SECONDS}
expected=$(awk '{print $2}' "${out_sshd_pub}")
while : ; do
	ssh-keyscan -t rsa "${host}" > "${out_keyscan}" 2>"${out_keyscan}.err"
	found=$(awk '{print $3}' "${out_keyscan}")
	[ -n "${found}" -a "${found}" = "${expected}" ] && break
	waited=$((${SECONDS}-${ssh_start}))
	[ $waited -lt $ssh_wait_max ] || {
		{
		echo "ssh keyscan failed, final ssh-keyscan output"
		echo "== output =="
		cat "${out_keyscan}"
		echo "== error =="
		cat "${out_keyscan}.err"
        echo "== expected =="
		echo "$expected"
		} | sed "s,^,$PREFIX," 1>&2
		fail "waited to long for ssh to get assigned key ($waited seconds)"
	}
	debug 2 "sleeping 10. $((${SECONDS}-${ssh_start}))/${ssh_wait_max} waited for ssh-keyscan"
	sleep 10
done

ssh_start=${SECONDS}
ssh_args=()
[ -n "${ssh_pubkey}" ] && ssh_args=( "${ssh_args[@]}" -i "${ssh_pubkey%.pub}" )

ssh_args=( -o UserKnownHostsFile=$known_hosts )

if [ -n "${known_hosts}" ]; then
	if [ -f "${known_hosts}" ]; then
		{
			{ [ -z "${associate_ip}" ] ||
			  ssh-keygen -R "${associate_ip}" -f "${known_hosts}"; } &&
			ssh-keygen -R "${host}" -f "${known_hosts}"
		} > "${out_clean_knownhosts}" 2>&1 ||
		{
			cat "${out_clean_knownhosts}";
			fail "failed to remove ${host}/${associate_ip} from ${known_hosts}"
		}
	fi
	cat "${out_keyscan}" >> "${known_hosts}" ||
		fail "failed to write to ${known_hosts}"
fi

while ! sshconnect "${ssh_args[@]}" "${SSH_USER}@${host}" 2>"$connect_err"; do
	waited=$((${SECONDS}-${ssh_start}))
	[ $waited -lt $ssh_wait_max ] || {
		error "ssh error:";
		cat "$connect_err" 1>&2;
		fail "waited to long for ssh ($waited seconds)"
	}
	debug 2 "sleeping 10. $((${SECONDS}-${ssh_start}))/${ssh_wait_max} waited for ssh-connect"
	sleep 10
done

debug 2 "ssh verified, total time $((${SECONDS}-${START_TIME}))s"

echo "${iid}" "${host}"
exit 0