~ubuntu-on-ec2/ubuntu-on-ec2/cloud-utils

#!/bin/bash
# vi: ts=4 noexpandtab
# This script uses bash arrays; do not switch to /bin/sh
#
#    uec-publish-image - wrapper for uec image publishing
#
#    Copyright (C) 2010 Canonical Ltd.
#
#    Authors: Scott Moser <smoser@canonical.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, version 3 of the License.
#
#    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, see <http://www.gnu.org/licenses/>.


EC2PRE=${EC2PRE:-euca-}
TMPD=""
RENAME_D=""
VERBOSITY=0
IMAGE_TYPES=( auto 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 bucket

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

   options:
      -l|--add-launch <user_id>   : user_id can be "all", or "none"
         --dry-run                : only report what would be done
         --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
        |--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
         --name   <name>          : register with '--name'.
                                    default: basename(publish_path)

   if type is 'image', then:
      -k | --kernel  k        : use previously registered kernel with id 'k'
                                specify 'none' for no kernel
      -K | --kernel-file f    : bundle, upload, use file 'f' as kernel
      -r | --ramdisk r        : use previously registered ramdisk with id 'r'
                                specify 'none' for no ramdisk
      -R | --ramdisk-file f   : bundle, upload, use file 'f' as ramdisk
      -B | --block-device-mapping m : specify block device mapping in bundle
EOF
}

bad_Usage() { Usage 1>&2; [ $# -eq 0 ] || error "$@"; exit 1; }
cleanup() {
	local x=""
	for x in "${RENAME_D}" "${TMPD}"; do
		[ -z "${x}" -o ! -d "${x}" ] || rm -Rf "${x}"
	done
	return 0
}

debug() {
	local level=${1}
	shift;
	[ "${level}" -ge "${VERBOSITY}" ] && return
	error "$(date):" "${@}"
}
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=$?
	echo "failed: ${*}"
	cat "${dir}/${pre}.stdout"
	cat "${dir}/${pre}.stderr" 1>&2
	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
}

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

get_manifest_id() {
	local str="${1}" cmp="~" tmpf="" out="" ret=1
	out=$(${EC2PRE}describe-images -a |
	        awk '$3 '$cmp' p { printf("%s\t%s\n",$2,$3); }' "p=${str}" ;
			checkstatus ${PIPESTATUS[@]}) || return 1
	_RET=${out}
	return
}
get_image_type() {
	local image=${1} file_out="" img_type=""
	file_out=$(file --uncompress "${image}") || return 1;
	case "${file_out}" in
		*[lL]inux\ kernel*) img_type="kernel";;
		*LSB\ executable*gzip*) img_type="kernel";;
		*cpio\ archive*) img_type="ramdisk";;
		*ext[234]\ file*|*boot\ sector*) img_type="image";;
		*) error "unable to determine image type. pass --type"; return 1;;
	esac
	_RET=${img_type}
	return 0
}

upload_register() {
	local out=""
	out=$(uec-publish-image "${@}") || return
	set -- ${out}
	_RET=${1}
}

short_opts="B:h:k:K:l:no:r:R:t:vw:"
long_opts="add-launch:,allow-existing,block-device-mapping:,dry-run,help,kernel:,kernel-file:,name:,output:,ramdisk:,ramdisk-file:,rename:,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=""
dry_run=0
image=""
img_type="auto"
kernel=""
kernel_file=""
output=""
ramdisk=""
ramdisk_file=""
rename=""
name=""
wdir_in=""
dev_mapping=""

while [ $# -ne 0 ]; do
	cur=${1}; next=${2};
	case "$cur" in
		-d|--working-dir) wdir_in=${next}; shift;;
		-h|--help) Usage; exit 0;;
		-B|--block-device-mapping) dev_mapping=${next}; shift;;
		-k|--kernel) kernel=${next}; shift;;
		-K|--kernel-file) kernel_file=${next}; 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;;
		-r|--ramdisk) ramdisk=${next}; shift;;
		-R|--ramdisk-file) ramdisk_file=${next}; shift;;
		-n|--dry-run) dry_run=1;;
		   --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;;
		--) shift; break;;
		-*) bad_Usage "confused by ${cur}";;
	esac
	shift;
done

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

# remove any trailing slashes on bucket
while [ "${bucket%/}" != "${bucket}" ]; do bucket=${bucket%/}; done

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

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

[ -f "${image}" ] || bad_Usage "${image}: image is not a file"

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

[ -n "${kernel_file}" -a -n "${kernel}" ] &&
	bad_Usage "--kernel-file is incompatible with --kernel"
[ -n "${ramdisk_file}" -a -n "${ramdisk}" ] &&
	bad_Usage "--ramdisk-file is incompatible with --ramdisk"

[ -z "${kernel_file}" -o -f "${kernel_file}" ] ||
	fail "${kernel_file} is not a file"
[ -z "${ramdisk_file}" -o -f "${ramdisk_file}" ] ||
	fail "${ramdisk_file} is not a file"

if [ "${img_type}" = "auto" ]; then
	get_image_type "${image}" ||
		fail "failed to determine file type of ${image}"
	img_type=${_RET}
fi

[ -n "${dev_mapping}" -a "${img_type}" != "image" ] &&
	fail "-B/--block-device-mapping can only be specified for --type=image"

[ -n "${rename}" ] || rename=${image##*/}

image_full=$(readlink -f "${image}") ||
	fail "failed to get full path to ${image}"

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}"
fi
trap cleanup EXIT

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

# bundle-kernel doesn't like for file to exist in destination-dir
# so, create it one dir under there
RENAME_D=$(mktemp -d "${wdir}/.rename.XXXXXX") &&
	ln -s "${image_full}" "${RENAME_D}/${rename}" &&
	rename_full="${RENAME_D}/${rename}" ||
	fail "link failed: working-dir/rename/${rename} -> ${image_full}"

reg_id=""

manifest="${rename}.manifest.xml"

# set up "pass through" args to go through to kernel/ramdisk publishing
pthr=( )
[ $VERBOSITY -eq 0 ] || pthr[${#pthr[@]}]="--verbose"
[ ${allow_existing} -eq 0 ] || pthr[${#pthr[@]}]="--allow-existing"
[ -z "${add_acl}" ] ||
	{ pthr[${#pthr[@]}]="--add-launch"; pthr[${#pthr[@]}]="${add_acl}"; }
[ ${dry_run} -eq 0 ] || pthr[${#pthr[@]}]="--dry-run"

if [ -n "${kernel_file}" ]; then
	debug 1 "publishing kernel ${kernel_file}"
	upload_register --type kernel "${pthr[@]}" \
		"${arch}" "${kernel_file}" "${bucket}" ||
		fail "failed to register ${kernel_file}"
	kernel=${_RET}
	debug 1 "kernel registered as ${kernel}"
fi

if [ -n "${ramdisk_file}" ]; then
	debug 1 "publishing ramdisk ${ramdisk_file}"
	upload_register --type ramdisk "${pthr[@]}" \
		"${arch}" "${ramdisk_file}" "${bucket}" ||
		fail "failed to register ${ramdisk_file}"
	ramdisk=${_RET}
	debug 1 "ramdisk registered as ${ramdisk}"
fi

if [ ${VERBOSITY} -ge 1 -o ${dry_run} -ne 0 ]; then
	[ -n "${kernel}" ] && krd_fmt=" %s/%s" &&
		krd_args=( "${kernel}" "${ramdisk:-none}" )
	errorp "[%-6s] %s => %s/%s ${krd_fmt}\n" "${img_type}" \
		"${image##*/}" "${bucket}" "${rename}" "${krd_args[@]}"
	if [ ${dry_run} -ne 0 ]; then
		case "${img_type}" in
			kernel) pre="eki";;
			ramdisk) pre="eri";;
			image) pre="emi";;
		esac
		printf "%s\t%s\n" "${pre}-xxxxxxxx" "${bucket}/${rename##*/}"
		exit
	fi
fi

krd_args=( );
[ -n "${kernel}" -a "${kernel}" != "none" ] &&
	krd_args=( "${krd_args[@]}" "--kernel" "${kernel}" )
[ -n "${ramdisk}" -a "${ramdisk}" != "none" ] &&
	krd_args=( "${krd_args[@]}" "--ramdisk" "${ramdisk}" )

if [ -z "${name}" ]; then
	name=${rename}
	debug 1 "using ${name} for --name"
fi

if [ "${EC2PRE%ec2-}" != "${EC2PRE}" ]; then
	req="EC2_CERT EC2_PRIVATE_KEY EC2_USER_ID EC2_ACCESS_KEY EC2_SECRET_KEY"
	for env_name in ${req}; do
		[ -n "${!env_name}" ] ||
			fail "when using ec2- tools, you must set env: ${req}"
	done
	ex_bundle_args=( --cert "${EC2_CERT}" 
	                 --privatekey "${EC2_PRIVATE_KEY}" 
	                 --user "${EC2_USER_ID}" )
	ex_upload_args=( --access-key "${EC2_ACCESS_KEY}" 
	                 --secret-key "${EC2_SECRET_KEY}" )

fi

debug 1 "checking for existing registered image at ${bucket}/${manifest}"
get_manifest_id "${bucket}/${manifest}" ||
	fail "failed to check for existing manifest"
if [ -n "${_RET}" ]; then
	set -- ${_RET}
	img_id=${1}; path=${2}
	[ ${allow_existing} -eq 1 ] ||
		fail "${bucket}/${manifest} already registered as ${img_id}"
	debug 1 "using existing ${img_id} for ${bucket}/${manifest}"
else
	bundle_args=( "--image" "${rename_full}" )
	[ -n "${dev_mapping}" ] &&
		bundle_args[${#bundle_args[@]}]="--block-device-mapping=${dev_mapping}"

	case "${img_type}" in
		kernel|ramdisk)
			bundle_args[${#bundle_args[@]}]="--${img_type}"
			bundle_args[${#bundle_args[@]}]="true"
	esac
	run "${wdir}" "bundle" "bundling ${img_type} ${image}" \
		${EC2PRE}bundle-image --destination "${wdir}" --arch "${arch}" \
			"${ex_bundle_args[@]}" \
			"${bundle_args[@]}" "${krd_args[@]}" ||
		fail "failed to bundle ${img_type} ${image}"

	run "${wdir}" "upload" "upload ${bucket}/${manifest}" \
		${EC2PRE}upload-bundle --bucket "${bucket}" \
				"${ex_upload_args[@]}" \
				--manifest "${wdir}/${manifest}" ||
			fail "failed to upload bundle to ${bucket}/${manifest}"

	junk="" img_id="";
	run "${wdir}" "register" "register ${bucket}/${manifest}" \
		${EC2PRE}register ${name:+--name "${name}"} \
			"${ex_register_args[@]}" "${bucket}/${manifest}" &&
		read junk img_id < "${wdir}/register.stdout" &&
		[ "${img_id#???-}" != "${img_id}" ] || {
			if bad=$(get_manifest_id "${bucket}/${manifest}") &&
			   [ -n "${bad}" ]; then
				set -- ${bad}
				bad_id=${1}
				error "un-registering invalid $bad" >/dev/null
				${EC2PRE}deregister "${bad_id}"
			fi
			fail "failed to register ${manifest}"
		}

	debug 1 "registered at ${bucket}/${manifest} as ${img_id}"

fi
debug 1 "${img_id} ${bucket}/${manifest}"

if [ -z "${output}" -o "${output}" = "-" ]; then
	printf "%s\t%s\n" "${img_id}" "${bucket}/${manifest}"
else
	printf "%s\t%s\n" "${img_id}" "${bucket}/${manifest}" >> "${output}"
fi

for user in ${add_acl}; do
	run "${wdir}" "add_user.${user}" \
		"add ${user} to ${manifest}" \
		${EC2PRE}modify-image-attribute \
			--launch-permission --add "${user}" "${img_id}" ||
		fail "failed to add launch permission for ${user} to ${img_id}"
done

exit 0