~josvaz/ubuntu-on-ec2/ec2-publishing-scripts-lp1613470-udev-cdrom

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

TMPD=""
RENAME_D=""
EXTRACT_D=""
VERBOSITY=0
REGIONS=( )
[ -n "${#EC2_ALL_REGIONS}" ] && REGIONS=( ${EC2_ALL_REGIONS} )

DEFAULT_MAPPER=append_full

error() { echo "$@" 1>&2; }
errorp() { printf "$@" 1>&2; }
fail() { [ $# -eq 0 ] || error "$@"; exit 1; }
failp() { [ $# -eq 0 ] || errorp "$@"; exit 1; }

Usage() {
	cat <<EOF
Usage: ${0##*/} [ options ] arch image

   arch           : one of i386 or x86_64
   image          : the image to upload and register

   options:
      -l|--add-launch <user_id>   : user_id can be "all", or "none"
         --bucket-mapper          : relavent only with bucket-prefix
                                    default: ${DEFAULT_MAPPER}
         --bucket-prefix <pre>    : bucket prefix to use
         --buckets       <list>   : list of buckets to publish to
         --dry-run                : only report what would be done
         --name          name     : register with '--name' to ec2-register
         --path-prefix   <pre>    : path inside bucket to put image in
         --resize        size     : resize the image with 'qemu-img resize'
                                    before publishing.
         --regions       <list>   : list of regions to publish to
                                    default: all (per describe-regions)
         --allow-existing         : if a image is already registered
                                    simply report as if work was done
		 --sriov				  : enable SRIOV support
      -o|--output <file>          : write registered id and manifest to file
      -r|--rename <publish_path>  : publish to bucket/<publish_path>
                                    default: bucket/<basename(image)>
      -v|--verbose                : increase verbosity

   all <list> entries are comma or space delimited, and map 1x1 with regions
EOF
}

bad_Usage() { Usage 1>&2; [ $# -eq 0 ] || error "$@"; exit 1; }
cleanup() {
	local pids
	pids=$(jobs -pr)
	if [ -n "$pids" ]; then
		error "waiting on pids: $pids"
		kill $pids
		wait $pids
	fi
	local x=""
	for x in "$EXTRACT_D" "${RENAME_D}" "${TMPD}"; do
		[ -z "${x}" -o ! -d "${x}" ] || rm -Rf "${x}"
	done

	return 0
}

debug() {
	local level=${1}
	shift;
	[ "${level}" -gt "${VERBOSITY}" ] && return
	echo "$(date): ${@}" 1>&2
}
run() {
	local dir="${1}" pre=${2} msg=${3};
	shift 3;
	[ -e "${dir}/stamp.${pre}" ] &&
		{ debug 1 "skipping ${pre}"; return 0; }
	debug 1 "${msg}"
	echo "$@" > "${dir}/${pre}.cmd"
	"$@" > "${dir}/${pre}.stdout" 2> "${dir}/${pre}.stderr" &&
		: > "${dir}/stamp.${pre}" && return 0
	local ret=$?
	cat "${dir}/${pre}.stderr" 1>&2
	return ${ret}
}

retry() {
	local trycount=${1} sleep=${2}
	shift; shift;
	local i=0 smsg=" sleeping ${sleep}: $*" ret=0
	for((i=0;i<${trycount};i++)); do
		"$@" && return 0
		ret=$?
		[ $(($i+1)) -eq ${trycount} ] && smsg=""
		debug 1 "Warning: cmd failed [try $(($i+1))/${trycount}].${smsg}"
		sleep $sleep
	done
	return $ret
}

search_args() {
	local x="" i=0 needle="$1"
	shift;
	for x in "${@}"; do
		[ "${needle}" = "${x}" ] && { _RET=$i; return 0; }
		i=$(($i+1))
	done
	return 1
}

append_location() {
	local pre=${1} region=${2} rsuf=""
	[ -n "${pre}" -a -n "${region}" ] || exit 1;
	region2loc "${region}" &&
		rsuf=-$(echo "${_RET}" | tr '[:upper:]' '[:lower:]') &&
		echo "${pre}${rsuf}"
}

append_short() {
	local pre=${1} region=${2}
	[ -n "${pre}" -a -n "${region}" ] || exit 1;
	local rsuf=-${region%%-*}
	# if the prefix ends with the suffix already, dont add it
	[ "${pre%${rsuf}}" != "${pre}" ] && { echo "${pre}"; return 0; }
	echo "${pre}${rsuf}"
}

append_full() {
	local pre=${1} region=${2}
	[ -n "${pre}" -a -n "${region}" ] || exit 1;
	# if the prefix ends with the region already, dont add it
	[ "${pre%${region}}" != "${pre}" ] && { echo "${pre}"; return 0; }
	echo "${pre}-${region}"
}

checkstatus() {
	local x="" i=0
	for x in "$@"; do
		[ "$x" = "0" ] || i=$(($i+1))
	done
	return $i
}

get_regions() {
	local regs=""
	[ "${#REGIONS[@]}" -eq 0 ] || return 0
	REGIONS=( $(ec2-list-all-regions) ) || return 1
}

region2loc() {
	[ "${1}" = "us-east-1" ] && { _RET="US"; return; }
	[ "${1}" = "eu-west-1" ] && { _RET="EU"; return; }
	_RET=${1}; return;
}


explode() {
	local delim="$1" str="${2}" limit="${3:-0}"
	_RET=( );
	local ostr=""
	while [ ${#_RET[@]} -lt ${limit} -o ${limit} -eq 0 ]; do
		ostr=${str}
		str=${str#*[${delim}]}
		[ ${limit} -eq $((${#_RET[@]}+1)) -o "${str}" = "${ostr}" ] &&
			{ _RET[${#_RET[@]}]=${ostr}; return 0; }
		_RET[${#_RET[@]}]=${ostr%?${str}}
	done
	return
}
split_input() {
	[ -n "${1}" ] || { _RET=( ); return 0; }
	explode ", " "${1}"
}

select_migrate_source() {
	# select_migrate_source(target, available_regions)
	local target=$1 i=""
	shift;
	#_RET=$1; return; # just return the first available
	_RET="" # never migrate
}

short_opts="hl:no:r:t:vw:"
long_opts="add-launch:,allow-existing,buckets:,bucket-mapper:,bucket-prefix:,dry-run,help,name:,output:,path-prefix:,regions:,rename:,resize:,type:,verbose,working-dir:,sriov"
getopt_out=$(getopt --name "${0##*/}" \
	--options "${short_opts}" --long "${long_opts}" -- "$@") &&
	eval set -- "${getopt_out}" ||
	bad_Usage

add_acl=""
allow_existing=0
arch=""
bucket_pre=""
buckets=( )
bucket_mapper=${DEFAULT_MAPPER}
dry_run=0
image=""
locations=( )
output=""
path_prefix=""
regions=( )
rename=""
resize=""
name=""
tries=3;
wdir_in=""
sriov=0

while [ $# -ne 0 ]; do
	cur=${1}; next=${2};
	case "$cur" in
		-d|--working-dir) wdir_in=${next}; shift;;
		-h|--help) Usage; exit 0;;
		-l|--add-launch)
			if [ "${next}" = "none" ]; then
				add_acl=""
			else
				user=${next//-/}; # just be nice and remove '-'
				add_acl="${add_acl:+${add_acl} }${user}";
			fi
			shift;;
		   --name) name=${next}; shift;;
		-o|--output) output="${next}"; shift;;
		-n|--dry-run) dry_run=1;;
		-r|--rename) rename=${next}; shift;;
		-v|--verbose) VERBOSITY=$((${VERBOSITY}+1));;
		--allow-existing) allow_existing=1;;
		--bucket-mapper) bucket_mapper=${next}; shift;;
		--bucket-prefix) bucket_pre=${next}; shift;;
		--buckets)
			split_input "${next}"
			buckets=( "${buckets[@]}" "${_RET[@]}" );
			shift;;
		--path-prefix) path_prefix=${next}; shift;;
		--regions)
			if [ "${next}" = "all" ]; then
				get_regions || fail "failed to get region list"
				regions=( "${REGIONS[@]}" );
			else
				split_input "${next}"
				regions=( "${regions[@]}" "${_RET[@]}" );
			fi
			shift;;
		--sriov) sriov=1;;
		--resize) resize=$next; shift;;
		--) shift; break;;
		-*) bad_Usage "confused by ${cur}";;
	esac
	shift;
done

[ $# -lt 2 ] && bad_Usage "must provide arch, image"
[ $# -gt 2 ] && bad_Usage "unexpected arguments: ${3}"
arch="${1}"
image="${2}"
image_in="$image"

# Make sure that non-amd64 builds don't happen
[ "${arch}" == "amd64" ] || {
	debug 1 "HVM instance-store publication is only supported on amd64, not ${arch}";
	exit 0;
}

[ "$1" = "-h" -o "$1" = "--help" ] && { Usage; exit 0; }

[ "${arch}" = "amd64" ] && arch=x86_64 && debarch=amd64 || debarch=${arch}

[ "$arch" = "i386" -o "$arch" = "x86_64" ] ||
	bad_Usage "arch must be i386 or x86_64"

[ -n "${bucket_pre}" -a "${#buckets[@]}" -ne 0 ] &&
	fail "bucket-prefix and bucket list are incompatible"

get_regions || fail "failed to get region list"

[ "${#regions[@]}" -eq 0 ] && regions=( "${REGIONS[@]}" )

for region in "${regions[@]}"; do
	search_args "${region}" "${REGIONS[@]}" ||
		bad_Usage "region ${region} is not valid"

	region2loc "${region}" ||
		bad_Usage "unable to map ${region} to a location"
	locations[${#locations[@]}]=${_RET}

	if [ -n "${bucket_pre}" ]; then
		bucket=$(${bucket_mapper} "${bucket_pre}" "${region}" "${arch}") ||
			fail "unable to map bucket name for ${region} (${bucket_mapper})"
		buckets[${#buckets[@]}]=${bucket}
	fi
done

[ "${#regions[@]}" -eq "${#buckets[@]}" ] ||
	failp "%s\n\t%s\n" "number of regions != number of buckets" \
		"see --buckets or --bucket-prefix"

[ -n "${rename}" ] || rename=${image##*/}
manifest_path="${path_prefix:+${path_prefix%/}/}${rename}.manifest.xml"

if [ -n "${wdir_in}" ]; then
	[ -d "${wdir_in}" ] || fail "input working directory not a directory";
	wdir=$(readlink -f "${wdir_in}") ||
		fail "failed to realize ${wdir_in}"
else
	TMPD=$(mktemp -d ${TMPDIR:-/tmp}/${0##*/}.XXXXXX) ||
		fail "failed to make tmpdir"
	wdir="${TMPD}"
	trap cleanup EXIT
fi

if [ -e "${wdir}/${rename}" ]; then
	[ "${wdir}/${rename}" -ef "${img}" ] ||
		fail "${wdir} already contains file named ${rename}"
fi

# if name is not given, then create one, use a shortened
# manifest name.  if bucket_pre is given, then use that as part
# of the name.
if [ -z "${name}" ]; then
	# there is no longer support for tools without '--name'
	name=${rename%.manifest.xml}; name=${name%.img};
	if [ -n "${bucket_pre}" ]; then
		name="${bucket_pre}/${path_prefix:+${path_prefix%/}/}${name}"
	fi
	register_name=( --name "${name}" )
fi

mf_bucket=""
to_publish=( )
registered_ids=( )
for((i=0;i<${#regions[@]};i++)); do
	region=${regions[$i]}
	bucket=${buckets[$i]}

	# hack! instance store and kernels renamed:
	# old 099720109477/ubuntu-images-milestone/ubuntu-lucid-10.04-beta1-amd64-server-20100317
	# new 099720109477/ubuntu/images-milestone/ubuntu-lucid-10.04-beta1-amd64-server-20100317
	# so we check for 'ubuntu.images' rather than what was existing
	exist_reg="/${name//ubuntu\//ubuntu.}$"
	exist_reg="/${name//hvm-/hvm.}$"
	debug 1 "checking for existing registered image $region at ${exist_reg}"
	tmp=$(xc2 ximages matching --region "$region" \
		"name ~ /${exist_reg//\//\\/}/ && name !~ /\/ebs\//" -o self) ||
		fail "failed to check for existing manifest in $region"

	if [ -n "${tmp}" ]; then
		tmp=$(echo "$tmp" | awk '-F\t' '{print $2, $3; exit(0); }')
		set -- ${tmp}
		img_id=${1}; path=${2}
		[ ${allow_existing} -eq 1 ] ||
			fail "$name already registered as ${img_id} ($path)"
		debug 1 "using existing ${img_id} for ${bucket}/${manifest_path}"

		# add this to a migrate from candidates
		mf_regions[${#mf_regions[@]}]=$region;
		registered_ids[$i]=${img_id}
	else
		to_publish[${#to_publish[@]}]=$i
        debug 1 "no existing registered image in ${regions[i]}"
	fi
done

need_image=0
for i in "${to_publish[@]}"; do
	region=${regions[$i]}
	select_migrate_source "$region" "${mf_regions[@]}"
	if [ -n "$_RET" ]; then
    	debug 1 "publish to $region, migrating from $_RET"
	else
		need_image=1
	fi
done

if [ $need_image -ne 0 ]; then
    pub_image="${TMPD}/${image##*/}"
	new_image="${pub_image//.img/.raw}"
    cp ${image} ${pub_image} ||
        fail "failed to copy image for manipulation"

	if [ -n "$resize" ]; then
		debug 1 "resizing image to $resize"
		qemu-img resize "${pub_image}" "$resize" 1>&2 ||
			fail "failed to resize ${image}"
	fi

  	debug 1 "converting qcow image to raw"
    qemu-img convert -O raw "${pub_image}" "${new_image}" ||
	  	fail "failed to convert image"
    image="${new_image}"

    if [ -n "${resize}" ]; then
        growpart "${image}" 1 ||
            fail "failed to grow partition images"
    fi

	RENAME_D=$(mktemp -d "${wdir}/.rename.XXXXXX") &&
		rename_full="${RENAME_D}/${rename}" &&
        mkdir -p $(dirname ${rename_full}) &&
		ln -sf "${image}" "${RENAME_D}/${rename}" ||
		fail "symlink failed: working-dir/rename/${rename} -> ${image_full}"

	bundle_cmd="bundle-image"
	bundle_args=( "--image" "${rename_full}" )

	run "${wdir}" "bundle.once" "bundling hvm-image ${image}" \
		xc2 "${bundle_cmd}" --destination "${wdir}" --arch "${arch}" \
			"${bundle_args[@]}" ||
		fail "failed to bundle hvm image"
fi

manifest="${rename##*/}.manifest.xml"
if [ "${#to_publish[@]}" -ne 0 ]; then
	for i in "${to_publish[@]}"; do
		region=${regions[$i]}
		location=${locations[$i]}
		bucket=${buckets[$i]}
		bucketpath="${bucket}${path_prefix:+/${path_prefix%/}}"

		# this would select a good source region for a migrate to region
		select_migrate_source "$region" "${mf_regions[@]}"
		if [ -n "$_RET" ]; then
			mf_region="$_RET"
			search_args "$mf_region" "${regions[@]}"
			mf_bucket=${buckets[$_RET]}
			[ -n "$mf_bucket" ] || fail "did not find $region in $regions"
		fi

		bucketpath="${bucket}${path_prefix:+/${path_prefix%/}}"
		( # start a subshell so we can run each region parallel

		if [ -z "$mf_bucket" ]; then
			debug 2 xc2 upload-bundle --bucket "${bucketpath}" \
					--manifest "${wdir}/${manifest}" --retry
			retry 20 1m \
				run "${wdir}" "upload.${region}" \
					"upload ${bucket}/${manifest_path} to ($region)" \
				xc2 upload-bundle --bucket "${bucketpath}" \
					--manifest "${wdir}/${manifest}" --retry || {
					error "failed to upload bundle to ${bucket}/${manifest_path}";
					exit 1;
					}
		else
			debug 2 xc2 migrate-image --bucket "${mf_bucket}" \
					--manifest "${manifest_path}" \
					--destination-bucket "${bucket}" \
					--region "${region}" ||
			retry 20 1m \
				run "${wdir}" "migrate.${region}" \
					"migrating HVM image  ${manifest_path} (${mf_bucket}=>${region})" \
				xc2 migrate-image --bucket "${mf_bucket}" \
					--manifest "${manifest_path}" \
					--destination-bucket "${bucket}" \
					--region "${region}" ||
					{ error "failed to migrate to ${region}"; exit 1; }
		fi
		junk="" img_id="" sriov_args="";
		[ "${sriov:-0}" -eq 1 ] && sriov_args="--sriov=simple";
		run "${wdir}" "register.${region}" "register ${bucket}/${manifest_path}" \
			xc2 register --region "${region}" \
				--architecture "${arch}" --virtualization-type=hvm \
				"${bucket}/${manifest_path}" \
				"${register_name[@]}" "${sriov_args}" && \
			read junk img_id < "${wdir}/register.${region}.stdout" &&
			[[ "${img_id}" =~ "ami-" ]] || {
				debug 1 $(<"${wdir}/register.${region}.stdout");
				errorp "failed to register ${bucket}/${manifest_path}.%s\n" \
                    "${junk:+ output:${junk} ${img_id}}";
				exit 1;
			}
		debug 1 "registered at ${region}: ${bucket}/${manifest_path} as ${img_id}"
		) &

		PUBLISHES[$i]=$!
	done

	for i in "${to_publish[@]}"; do
		region=${regions[$i]}
		wait "${PUBLISHES[$i]}" || fail "failed to publish to ${regions[$i]}"
		read junk img_id < "$wdir/register.$region.stdout"
		registered_ids[$i]=$img_id
	done
fi

label_type="hvm-instance"
for((i=0;i<${#regions[@]};i++)); do
	img_id=${registered_ids[$i]}
	bucket=${buckets[$i]}
	region=${regions[$i]}
	if [ -z "${output}" -o "${output}" = "-" ]; then
		printf "%s\t%s\t%s\t%s\t%s\n" "${region}" "${img_id}" "${debarch}" "${label_type}" "${bucket}/${manifest_path}"
	else
		printf "%s\t%s\t%s\t%s\t%s\n" "${region}" "${img_id}" "${debarch}" "${label_type}" "${bucket}/${manifest_path}" >> "${output}"
	fi
done

# now add acls for everything
for((i=0;i<${#registered_ids[@]};i++)); do
	img_id=${registered_ids[$i]}
	region=${reg_regions[$i]}
	for user in ${add_acl}; do
		run "${wdir}" "add_user.${region}.${user}" \
			"add ${user} to ${manifest} in ${region}" \
			xc2 modify-image-attribute --region "${region}" \
				--launch-permission --add "${user}" "${img_id}" ||
			fail "failed to add launch permission for ${user} to ${img_id}"
	done
done

exit 0