#!/bin/bash
# This script uses bash arrays; do not switch to /bin/sh
#
#    cloud-publish-image - wrapper for cloud 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 ID          : ID can be "all", or "none", or numerical ID
         --dry-run                : only report what would be done
         --allow-existing         : if a image is already registered
                                    simply report as if work was done
         --hook-img EXE           : invoke executable 'EXE' with full path to
                                    downloaded disk image file
      -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: publish_path

         --save-downloaded        : if the image is a url, save it to '.'

   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 tmpf="" out="" ret=1 m1="${1}" m2="${2}"
	out=$(${EC2PRE}describe-images -o self |
		awk '$3 ~ m1 || $3 ~ m2 { printf("%s\t%s\n",$2,$3); }' \
				"m1=$m1" "m2=${m2:-^$}"
			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=$(cloud-publish-image "${@}") || return
	set -- ${out}
	_RET=${1}
}

dl() {
	# dl url, target, quiet
	local url=${1} target=${2} quiet=${3:-1}
	if [ -f "${url}" ]; then
		[ "${target}" = "-" ] && { cat "$url"; return; }
		cp "$url" "$target"
		return
	fi
	local qflag="-q"
	[ "$quiet" = "0" ] && qflag=""

	wget $qflag --progress=dot:mega "$url" -O "$target" ||
		return 1
}

dl_input_image() {
	# this downloads an image if necessary and sets _RET to location of image
	local input="$1" save_dir="${2:-.}" ret="" quiet=1
	[ $VERBOSITY -ge 2 ] && quiet=0
	case "$input" in
		file://*)
			ret="$save_dir/${input##*/}"
			dl "${input#file://}" "$ret" $quiet || return $?;;
		http://*|ftp://*|https://*)
			ret="$save_dir/${input##*/}"
			dl "$input" "$ret" $quiet || return $?
			;;
		*) ret="$input";;
	esac
	_RET="$ret"
}


[ "${CLOUD_UTILS_WARN_UEC:-0}" = "0" ] && _n="${0##*/}" &&
	[ "${_n#uec}" != "${_n}" ] && export CLOUD_UTILS_WARN_UEC=1 &&
	error "WARNING: '${0##*/}' is now to 'cloud${_n#uec}'. Please update your tools or docs"

short_opts="B:h:k:K:l:no:r:R:t:vw:"
long_opts="add-launch:,allow-existing,block-device-mapping:,dry-run,help,hook-img:,kernel:,kernel-file:,name:,output:,image-to-raw,ramdisk:,ramdisk-file:,rename:,save-downloaded,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="image"
kernel=""
kernel_file=""
output=""
ramdisk=""
ramdisk_file=""
rename=""
save_dl=0
name=__unset__
wdir_in=""
dev_mapping=""
image2raw=0
raw_image=""
hook_img=""

while [ $# -ne 0 ]; do
	cur=${1}; next=${2};
	case "$cur" in
		-d|--working-dir) wdir_in=${next}; shift;;
		-h|--help) Usage; exit 0;;
		   --hook-img)
			[ -z "${hook_img}" ] || bad_Usage "only one --hook-img supported";
			[ -x "$next" ] || bad_Usage "--hook-img is not executable"
			hook_img=$(readlink -f "$next") ||
				bad_Usage "could not find full path to $next"
			hook_img="$next"
			shift;;
		-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;;
		   --image-to-raw) image2raw=1;;
		-r|--ramdisk) ramdisk=${next}; shift;;
		-R|--ramdisk-file) ramdisk_file=${next}; shift;;
		-n|--dry-run) dry_run=1;;
		   --rename) rename=${next}; shift;;
		   --save-downloaded) save_dl=1;;
		-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

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

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 [ -n "$kernel" -a "$kernel" != "none" ]; then
	aki_arch=""; ari_arch="";
	# if kernel is given, check that its arch matches the register arch
	aki_arch=""; ari_arch="";
	
	[ "$ramdisk" = "none" ] && _ramdisk="" || _ramdisk="$ramdisk"

	${EC2PRE}describe-images "$kernel" $_ramdisk > "${TMPD}/kernel.info" ||
		fail "failed to describe kernel ${kernel}"
	aki_arch=$(awk '-F\t' '$1 == "IMAGE" && $2 == id { print $8 }' \
		"id=$kernel" "$TMPD/kernel.info") && [ -n "$aki_arch" ] ||
		fail "failed to get arch of $kernel"
	if [ -n "$ramdisk" -a "$ramdisk" != "none" ]; then
		ari_arch=$(awk '-F\t' '$1 == "IMAGE" && $2 == id { print $8 }' \
			"id=$ramdisk" "$TMPD/kernel.info") && [ -n "$ari_arch" ] ||
			fail "failed to get arch of $ramdisk"
	fi

	# if kernel and ramdisk are given, and arch=i386 kernel/ramdisk=x86_64,
	# then assume loader kernel.
	case "$arch:$aki_arch:$ari_arch" in
		$arch:$arch:$arch|$arch:$arch:) : ;;
		i386:x86_64:x86_64)
			error "WARNING: assuming loader kernel ($kernel/$ramdisk arch=$aki_arch, provided arch=$arch)"
			arch="x86_64";;
		*) fail "arch $arch != kernel/ramdisk arch [$aki_arch/$ari_arch]";;
	esac
fi

save_dir="${wdir}"
[ $save_dl -eq 1 ] && save_dir=.

dl_input_image "$image" "$save_dir" && image="$_RET" ||
	fail "failed to download image $image to $save_dir"

[ -z "$kernel_file" ] ||
	{ dl_input_image "$kernel_file" "$save_dir" && kernel_file="$_RET"; } ||
	fail "failed to download kernel $kernel_file to $save_dir"

[ -z "$ramdisk_file" ] ||
	{ dl_input_image "$ramdisk_file" "$save_dir" && ramdisk_file="$_RET"; } ||
	fail "failed to download ramdisk $ramdisk_file to $save_dir"

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

[ -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##*/}

if [ "${name}" = "__unset__" ]; then

	# if user did not pass --name, try to figure out if register supports it
	# we unfortunately can't assume that '--help' exits 0
	${EC2PRE}register --help > "${TMPD}/register-help.out" 2>&1
	if grep -q -- "--name" "${TMPD}/register-help.out"; then
		name="${bucket}/${rename}"
		debug 1 "using ${name} for --name"
	else
		debug 1 "${EC2PRE}register seems not to support --name, not passing"
		name=""
	fi
	
elif [ -z "${name}" -o "${name}" == "none" ]; then
	# if user passed in '--name=""' or '--name=none", do not pass --name
	name=""
fi

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

if [ -e "${wdir}/${rename}" ]; then
	[ "${wdir}/${rename}" -ef "${image}" ] ||
		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 [ "${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}" "/$name$" ||
	fail "failed to check for existing manifest"
if [ -n "${_RET}" ]; then
	set -- ${_RET}
	img_id=${1}; path=${2}
	[ ${allow_existing} -eq 1 ] ||
		fail "${path} already registered as ${img_id}"
	debug 1 "using existing ${img_id} for ${bucket}/${manifest}"
else
	if [ $image2raw -eq 1 -a "$img_type" = "image" ]; then
		# this is really here because of LP: #836759 
		# but could be useful elsewhere
		qemu-img info "$image" > "${TMPD}/disk-info.out" ||
			fail "failed to qemu-img info $image"
		imgfmt=$(awk '-F:' '$1 == "file format" { sub(/ /,"",$2); print $2 }' \
			"${TMPD}/disk-info.out")
		if [ "$imgfmt" != "raw" ]; then
			debug 1 "converting image to raw"
			raw_image="${TMPD}/image.raw"
			qemu-img convert -O raw "$image" "$raw_image" ||
				fail "failed to convert image to raw"
			image="$raw_image"
			ln -sf "$raw_image" "$rename_full" ||
				fail "symlink to raw image $raw_image failed"
		else
			debug 1 "disk is already raw format, not converting"
		fi
	fi

	if [ -n "$hook_img" ]; then
		debug 1 "image hook: $hook_img $rename_full"
		"$hook_img" "$rename_full" ||
			fail "image hook failed: $hook_img ${rename_full} failed"
	fi

	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}" "/${name}") &&
			   [ -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

# vi: ts=4 noexpandtab