#!/bin/sh
#
#    cloud-resize-image - resize a cloud image
#
#    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/>.


Usage() {
	cat <<EOF
Usage: ${0##*/} [ options ] image size [output]
   Resize a partition image and contained filesystem to a new size.
   if output is given, do not modify 'image', but create new file 'output'

   New size is specified per resize2fs(8), e.g. "1G" for 1 gigabyte

   options:
      -v | --verbose    show command output
EOF
	return 0
}

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

human2bytes() {
	# converts size suitable for input to resize2fs to bytes
	# s:512 byte sectors, K:kilobytes, M:megabytes, G:gigabytes
	# none: block size of the image
	local input=${1} defunit=${2:-1024}
	local unit count;
	case "$input" in
		*s) count=${input%s}; unit=512;;
		*K) count=${input%K}; unit=1024;;
		*M) count=${input%M}; unit=$((1024*1024));;
		*G) count=${input%G}; unit=$((1024*1024*1024));;
		*)  count=${input}  ; unit=${2:-1024};;
	esac
	_RET=$((${count}*${unit}))
}

xtruncate() {
	if which truncate >/dev/null 2>&1; then
		truncate "${@}"
	else
		local size=${1} file=${2} blk=""
		size=${size#--size=}
		# this is a poor mans truncate supporting whatever human2bytes supports
		human2bytes "${size}" && blk=$((${_RET}/512)) &&
			dd if=/dev/zero of="${file}" obs=512 seek=${blk} count=0 2>/dev/null
	fi
}

runcmd() {
	local output=$1
	shift;
	if [ "$output" = "0" ]; then
		local out="" ret=0;
		out=$("${@}" 2>&1) || { ret=$?; error "${out}"; return $ret; }
	else
		"$@"
	fi
}

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

verbose=0
[ "$1" = "-v" -o "$1" = "--verbose" ] &&
	{ verbose=1; shift; }

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

[ $# -eq 3 -o $# -eq 2 ] || { Usage 1>&2; exit 1; }

old="$1"
size="$2"
new="${3:-${old}}"

[ -f "${old}" ] || fail "${old}: does not exist"

human2bytes "${size}" && new_size=${_RET} ||
	fail "failed to understand ${size}"

if [ ! "${old}" -ef "${new}" ]; then
	file_out=$(file "${old}") || fail "failed to read ${old} with 'file'"
	case "${file_out}" in
		*gzip\ compressed*)
			file_out_z=$(file -z "${old}")
			case "${file_out_z}" in
				*tar\ archive*)
					: > "${new}" && newd=$(dirname "${new}") ||
						fail "failed to get full path for ${new}"
					tmpd=$(mktemp -d "${newd}/.${0##*/}.XXXXXX") &&
						( cd "${tmpd}" && tar -S --wildcards -xzf - "*.img" &&
						  mv *.img "../${new}" ) < "${old}" || { 
						  rm -Rf "${tmpd}";
						  fail "failed to extract image from ${old}"
						}
					rm -Rf "${tmpd}"
					;;
				*)
					zcat -f "$old" | cp --sparse=always /dev/stdin "$new";;
			esac
			;;
		*) cp --sparse=always "${old}" "${new}";;
	esac
	[ $? -eq 0 ] || fail "failed to cp ${old} -> ${new}"
else
	# if old=new (in place), it must be a simple image file
	case "${old}" in
		*.gz) fail "refusing work in place compressed or archive file: ${old}";;
	esac
fi

ls_out=$(ls -l "${new}") &&
	old_size=$(echo "${ls_out}" | awk '{print $5}') ||
	fail "failed to get size of ${new_img}"

runcmd "${verbose}" e2fsck -fp "${new}" ||
	fail "failed to fsck ${new}"

if [ "${old_size}" -lt "${new_size}" ]; then
	xtruncate "--size=$size" "$new" || fail "failed to change size of ${new}"
fi

runcmd "${verbose}" resize2fs "$new" "$size" ||
	fail "failed to resize ${new} -> ${size}"

if [ "${old_size}" -gt "${new_size}" ]; then
	xtruncate "--size=$size" "$new" || fail "failed to change size of ${new}"
fi

echo "resized ${new} to ${size}"

exit 0

# vi: ts=4 noexpandtab