~philroche/ubuntu-on-ec2/ec2-publishing-scripts-cloud-images-rsync-module

#!/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
IMAGE_TYPES=( image kernel ramdisk vmlinuz initrd )

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 'resize-cloud-image'
                                    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
      -o|--output <file>          : write registered id and manifest to file
      -r|--rename <publish_path>  : publish to bucket/<publish_path>
                                    default: bucket/<basename(image)>
      -t|--type   <type>          : type is one of kernel/ramdisk/image
      -v|--verbose                : increase verbosity

   if type is 'image', then:
      -i|--ramdisks  <list>       : the aris to register (list maps to regions)
      -k|--kernels   <list>       : the akis to register (list maps to regions)

   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="hi:k:l:no:r:t:vw:"
long_opts="add-launch:,allow-existing,buckets:,bucket-mapper:,bucket-prefix:,dry-run,help,kernels:,name:,output:,path-prefix:,ramdisks:,regions:,rename:,resize:,type:,verbose,working-dir:"
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=""
img_type="image"
kernel=""
kernels=( )
locations=( )
output=""
path_prefix=""
ramdisk=""
ramdisks=( )
regions=( )
rename=""
resize=""
name=""
tries=3;
wdir_in=""

while [ $# -ne 0 ]; do
	cur=${1}; next=${2};
	case "$cur" in
		-d|--working-dir) wdir_in=${next}; shift;;
		-h|--help) Usage; exit 0;;
		-k|--kernels)
			split_input "${next}"
			kernels=( "${kernels[@]}" "${_RET[@]}" );
			shift;;
		-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;;
		-i|--ramdisks)
			split_input "${next}"
			ramdisks=( "${ramdisks[@]}" "${_RET[@]}" );
			shift;;
		-n|--dry-run) dry_run=1;;
		-r|--rename) rename=${next}; shift;;
		-t|--type) 
			img_type=${next};
			search_args "${img_type}" "${IMAGE_TYPES[@]}" ||
				bad_Usage "image type (${next}) not in ${IMAGE_TYPES[*]}"
			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;;
		--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"

[ "$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

[ ${#kernels[@]} -eq 0 -o ${#kernels[@]} -eq ${#regions[@]} ] ||
	fail "must provide kernel for each region"

[ ${#ramdisks[@]} -eq 0 -o ${#ramdisks[@]} -eq ${#regions[@]} ] ||
	fail "must provide ramdisk for each region"

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

[ "${img_type}" = "vmlinuz" ] && img_type="kernel"
[ "${img_type}" = "initrd" ] && img_type="ramdisk"

label_type=""
case "${img_type}" in
	kernel|ramdisk) label_type="${img_type}";;
	image) label_type="instance-store";;
	*) fail "unknown label type for ${img_type}";;
esac

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

if [ ${VERBOSITY} -ge 1 -o ${dry_run} -ne 0 ]; then
	for((i=0;i<${#regions[@]};i++)); do 
		krd_fmt=""; krd_args=( );
		[ -n "${kernels[$i]}" ] && krd_fmt=" %s/%s" &&
			krd_args=( "${kernels[$i]}" "${ramdisks[$i]}" )
		errorp "[%-6s] %s => %s/%s [%s]${krd_fmt}\n" "${img_type}" \
			"${image##*/}" "${buckets[$i]}${path_prefix:+/${path_prefix}}" \
			"${rename}" \
			"${regions[$i]}" "${krd_args[@]}"
	done
	[ ${dry_run} -ne 0 ] && exit
fi

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
	[ "${img_type}" = "image" ] ||
		failp "%s\n%s\n" \
			"for kernel/ramdisk, working-dir/<rename> cannot exist" \
			"(bug/annoyance in ec2-bundle-{kernel,ramdisk})"
	[ "${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.}$"
	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
	fi
done

# if any of the regions were not going to be migrate, then
# migrate-image only works for image, not kernel/initrd (LP: #428692)
need_image=0
for i in "${to_publish[@]}"; do
	region=${regions[$i]}
	select_migrate_source "$region" "${mf_regions[@]}"
	if [ -n "$_RET" -a "$img_type" = "image" ]; then
		debug 1 "publish to $region, migrating from $_RET"
	else
		need_image=1
		debug 1 "publish to $region"
	fi
done

if [ $need_image -ne 0 ]; then

	case "$image" in
		*.tar.gz)
			EXTRACT_D=$(mktemp -d "$wdir/.extract.XXXXXX") &&
			( cd "$EXTRACT_D" &&
				run "$wdir" extract "extracting tarball input" \
					tar -Sxvzf - ) < "$image" ||
			fail "failed to extract tarball input"
			# make sure there is 1 and only 1 image
			new_image=$(find "${EXTRACT_D}" -type f -name "*.img") &&
				[ -f "$new_image" ] ||
				fail "failed to find the image in tarball input $image"
			image=${new_image}
	esac

	if [ -n "$resize" ]; then
		# resize the image to 10G
		debug 1 "resizing image to $resize"
		resize-cloud-image "${image}" "$resize" 1>&2 ||
			fail "failed to resize ${image}"
	fi

	# bundle-kernel doesn't like for file to exist in destination-dir
	# so, create it one dir under there
	image_full=$(readlink -f "${image}") ||
		fail "failed to get full path to ${image}"

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

	bundle_cmd="bundle-${img_type}"
	case "${img_type}" in
		kernel) bundle_args=( "--kernel" "${rename_full}" ) ;;
		ramdisk) bundle_args=( "--image" "${rename_full}" ) ;;
		image) bundle_args=( "--image" "${rename_full}" ) ;;
		*) bad_Usage "type must be one of kernel, ramdisk, or image";;
	esac

	run "${wdir}" "bundle.once" "bundling ${img_type} ${image}" \
		xc2 "${bundle_cmd}" --destination "${wdir}" --arch "${arch}" \
			"${bundle_args[@]}" ||
		fail "failed to bundle ${img_type} 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]}
		kernel=${kernels[$i]}
		ramdisk=${ramdisks[$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

		krd_args=( ${kernel:+"--kernel=$kernel"}
			       ${ramdisk:+"--ramdisk=$ramdisk"} )

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

		retry 10 10s \
			xc2 create-bucket "$bucket" "$location" >/dev/null ||
			{ error "failed to create bucket"; exit 1; }

		# note, as found in LP: #428692 comment 10, migrate-image doesnt work
		# for kernel/initrd. So, use publish
		if [ -z "$mf_bucket" -o "$img_type" != "image" ]; 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 ${img_type} ${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="";
		run "${wdir}" "register.${region}" "register ${bucket}/${manifest_path}" \
			xc2 ximages register --region "${region}" "${bucket}/${manifest_path}" \
				"${register_name[@]}" \
				"${krd_args[@]}" && \
			read junk img_id < "${wdir}/register.${region}.stdout" &&
			[ "${img_id#???-}" != "${img_id}" ] || {
				# Warning: register with an invalid kernel or ramdisk
				# can register the image (per describe-images) even though
				# command ec2-register command will fail.
				# also, it does not write the new ami to stdout
				if bad=$(xc2 get-manifest-id --region "$region" \
					"${bucket}/${manifest_path}") &&
				   [ -n "${bad}" ]; then
					set -- ${bad}
					bad_id=${1}
					error "un-registering invalid $bad"
					xc2 deregister --region "$region" "${bad_id}"
				fi
				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

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}"
	else
		printf "%s\t%s\t%s\t%s\t%s\n" "${region}" "${img_id}" "${debarch}" "${label_type}" "${bucket}/${manifest}" >> "${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 euca 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