#!/bin/sh
#    Copyright (C) 2011 Canonical Ltd.
#    Copyright (C) 2013 Hewlett-Packard Development Company, L.P.
#
#    Authors: Scott Moser <smoser@canonical.com>
#             Juerg Haefliger <juerg.haefliger@hp.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/>.

# the fudge factor. if its within this many 512 byte sectors, dont bother
FUDGE=${GROWPART_FUDGE:-$((20*1024))}
TEMP_D=""
RESTORE_FUNC=""
RESTORE_HUMAN=""
VERBOSITY=0
DISK=""
PART=""
DRY_RUN=0

MBR_CHS=""
MBR_BACKUP=""
GPT_BACKUP=""

error() {
	echo "$@" 1>&2
}

fail() {
	[ $# -eq 0 ] || echo "FAILED:" "$@"
	exit 2
}

nochange() {
	echo "NOCHANGE:" "$@"
	exit 1
}

changed() {
	echo "CHANGED:" "$@"
	exit 0
}

change() {
	echo "CHANGE:" "$@"
	exit 0
}

cleanup() {
	if [ -n "${RESTORE_FUNC}" ]; then
		error "***** WARNING: Resize failed, attempting to revert" \
			"******"
		if ${RESTORE_FUNC} ; then
			error "***** Appears to have gone OK ****"
		else
			error "***** FAILED! or original partition table" \
				"looked like: ****"
			cat "${RESTORE_HUMAN}" 1>&2
		fi
	fi
	[ -z "${TEMP_D}" -o ! -d "${TEMP_D}" ] || rm -Rf "${TEMP_D}"
}

debug() {
	local level=${1}
	shift
	[ "${level}" -gt "${VERBOSITY}" ] && return
	error "${@}"
}

mktemp_d() {
	# just a mktemp -d that doens't need mktemp if its not there.
	_RET=$(mktemp -d "${TMPDIR:-/tmp}/${0##*/}.XXXXXX" 2>/dev/null) && \
		return
	_RET=$(umask 077 && t="${TMPDIR:-/tmp}/${0##*/}.$$" && \
		mkdir "${t}" &&	echo "${t}")
	return
}

Usage() {
	cat <<EOF
${0##*/} disk partition
   rewrite partition table so that partition takes up all the space it can
   options:
    -h | --help      print Usage and exit
         --fudge F   if part could be resized, but change would be
                     less than 'F', do not resize (default: ${FUDGE})
    -N | --dry-run   only report what would be done, show new 'sfdisk -d'
    -v | --verbose   increase verbosity / debug

   Example:
    - ${0##*/} /dev/sda 1
      Resize partition 1 on /dev/sda
EOF
}

bad_Usage() {
	Usage 1>&2
	error "$@"
	exit 2
}

mbr_restore() {
	sfdisk --no-reread "${DISK}" "${MBR_CHS}" -I "${MBR_BACKUP}"
}

mbr_resize() {
	RESTORE_HUMAN="${TEMP_D}/recovery"
	MBR_BACKUP="${TEMP_D}/orig.save"

	local change_out=${TEMP_D}/change.out
	local dump_out=${TEMP_D}/dump.out
	local new_out=${TEMP_D}/new.out
	local dump_mod=${TEMP_D}/dump.mod
	local tmp="${TEMP_D}/tmp.out"
	local err="${TEMP_D}/err.out"

	local _devc cyl _word1 heads _word2 sectors _word3 tot dpart
	local pt_start pt_size pt_end max_end new_size change_info

	# --show-pt-geometry outputs something like
	#     /dev/sda: 164352 cylinders, 4 heads, 32 sectors/track
	sfdisk "${DISK}" --show-pt-geometry > "${tmp}" 2>"${err}" && \
		read _devc cyl _word1 heads _word2 sectors _word3 < \
		"${tmp}" && MBR_CHS="-C ${cyl} -H ${heads} -S ${sectors}" || \
		fail "failed to get CHS from ${DISK}"

	tot=$((${cyl}*${heads}*${sectors}))

	debug 1 "geometry is ${MBR_CHS}. total size=${tot}"
	sfdisk "${MBR_CHS}" -uS -d "${DISK}" > "${dump_out}" 2>"${err}" || \
		fail "failed to dump sfdisk info for ${DISK}"

	{
		echo "## sfdisk ${MBR_CHS} -uS -d ${DISK}"
		cat "${dump_out}"
	}  > "${RESTORE_HUMAN}"
	[ $? -eq 0 ] || fail "failed to save sfdisk -d output"

	[ ${VERBOSITY} -gt 0 ] && cat ${RESTORE_HUMAN}

	sed -e 's/,//g; s/start=/start /; s/size=/size /' "${dump_out}" > \
		"${dump_mod}"

	dpart="${DISK}${PART}" # disk and partition number
	if [ -b "${DISK}p${PART}" -a "${DISK%[0-9]}" != "${DISK}" ]; then
		# for block devices that end in a number (/dev/nbd0)
		# the partition is "<name>p<partition_number>" (/dev/nbd0p1)
		dpart="${DISK}p${PART}"
	elif [ "${DISK#/dev/loop[0-9]}" != "${DISK}" ]; then
		# for /dev/loop devices, sfdisk output will be <name>p<number>
		# format also, even though there is not a device there.
		dpart="${DISK}p${PART}"
	fi

	pt_start=$(awk '$1 == pt { print $4 }' "pt=${dpart}" < \
		"${dump_mod}") && \
		pt_size=$(awk '$1 == pt { print $6 }' "pt=${dpart}" < \
		"${dump_mod}") && \
		[ -n "${pt_start}" -a -n "${pt_size}" ] && \
		pt_end=$((${pt_size}+${pt_start})) || \
		fail "failed to get start and end for ${dpart} in ${DISK}"

	# find the minimal starting location that is >= pt_end 
	max_end=$(awk '$3 == "start" { if($4 >= pt_end && $4 < min) \
		{ min = $4 } } END { printf("%s\n",min); }' \
		min=${tot} pt_end=${pt_end} "${dump_mod}") && \
		[ -n "${max_end}" ] || \
		fail "failed to get max_end for partition ${PART}"

	debug 1 "max_end=${max_end} tot=${tot} pt_end=${pt_end}" \
		"pt_start=${pt_start} pt_size=${pt_size}"
	[ $((${pt_end})) -eq ${max_end} ] && \
		nochange "partition ${PART} is size ${pt_size}. it cannot be" \
		"grown"
	[ $((${pt_end}+${FUDGE})) -gt ${max_end} ] && \
		nochange "partition ${PART} could only be grown by" \
		"$((${max_end}-${pt_end})) [fudge=${FUDGE}]"

	# now, change the size for this partition in ${dump_out} to be the
	# new size
	new_size=$((${max_end}-${pt_start}))
	sed "\|^\s*${dpart} |s/${pt_size},/${new_size},/" "${dump_out}" > \
		"${new_out}" ||	fail "failed to change size in output"

	change_info="partition=${PART} start=${pt_start} old: size=${pt_size} end=${pt_end} new: size=${new_size},end=${max_end}"
	if [ ${DRY_RUN} -ne 0 ]; then
		echo "CHANGE: ${change_info}"
		{
			echo "# === old sfdisk -d ==="
			cat "${dump_out}"
			echo "# === new sfdisk -d ==="
			cat "${new_out}"
		} 1>&2
		exit 0
	fi

	sfdisk --no-reread "${DISK}" "${MBR_CHS}" --force \
		-O "${MBR_BACKUP}" < "${new_out}" > "${change_out}" 2>&1 || {
		RESTORE_FUNC=mbr_restore
		error "attempt to resize ${DISK} failed. sfdisk output below:"
		sed 's,^,| ,' "${change_out}" 1>&2
		fail "failed to resize"
	}
	changed "${change_info}"

	# dump_out looks something like:
	## partition table of /tmp/out.img
	#unit: sectors
	#
	#/tmp/out.img1 : start=        1, size=    48194, Id=83
	#/tmp/out.img2 : start=    48195, size=   963900, Id=83
	#/tmp/out.img3 : start=  1012095, size=   305235, Id=82
	#/tmp/out.img4 : start=  1317330, size=   771120, Id= 5
	#/tmp/out.img5 : start=  1317331, size=   642599, Id=83
	#/tmp/out.img6 : start=  1959931, size=    48194, Id=83
	#/tmp/out.img7 : start=  2008126, size=    80324, Id=83
}

gpt_restore() {
	sgdisk -l "${GPT_BACKUP}" "${DISK}"
}

gpt_resize() {
	GPT_BACKUP="${TEMP_D}/pt.backup"

	local pt_info="${TEMP_D}/pt.info"
	local pt_pretend="${TEMP_D}/pt.pretend"
	local pt_data="${TEMP_D}/pt.data"
	local out="${TEMP_D}/out"

	local dev="disk=${DISK} partition=${PART}"

	local pt_start pt_end pt_size last pt_max code guid name new_size
	local old new change_info
	
	# Dump the original partition information and details to disk. This is
	# used in case something goes wrong and human interaction is required
	# to revert any changes.
	sgdisk -i "${PART}" -p "${DISK}" > "${pt_info}" || \
		fail "${dev}: failed to dump original sgdisk info"
	RESTORE_HUMAN="${pt_info}"

	[ "${VERBOSITY}" -gt 0 ] && {
		echo "${dev}: original sgdisk info:"
		cat "${pt_info}"
	}

	# Pretend to move the backup GPT header to the end of the disk and dump
	# the resulting partition information. We use this info to determine if
	# we have to resize the partition.
	sgdisk -P -e -p "${DISK}" > "${pt_pretend}" || \
		fail "${dev}: failed to dump pretend sgdisk info"

	[ "${VERBOSITY}" -gt 0 ] && {
		echo "${dev}: pretend sgdisk info:"
		cat "${pt_pretend}"
	}

	# Extract the partition data from the pretend dump
	awk 'found { print } ; $1 == "Number" { found = 1 }' \
		"${pt_pretend}" > "${pt_data}" || \
		fail "${dev}: failed to parse pretend sgdisk info"

	# Get the start and end sectors of the partition to be grown
	pt_start=$(awk '$1 == '"${PART}"' { print $2 }' "${pt_data}") && \
		[ -n "${pt_start}" ] || \
		fail "${dev}: failed to get start sector"
	pt_end=$(awk '$1 == '"${PART}"' { print $3 }' "${pt_data}") && \
		[ -n "${pt_end}" ] || \
		fail "${dev}: failed to get end sector"
	pt_size="$((${pt_end} - ${pt_start}))"

	# Get the last usable sector
	last=$(awk '/last usable sector is/ { print $NF }' \
		"${pt_pretend}") && [ -n "${last}" ] || \
		fail "${dev}: failed to get last usable sector"

	# Find the minimal start sector that is >= pt_end 
	pt_max=$(awk '{ if ($2 >= pt_end && $2 < min) { min = $2 } } END \
		{ print min }' min="${last}" pt_end="${pt_end}" \
		"${pt_data}") && [ -n "${pt_max}" ] || \
		fail "${dev}: failed to find max end sector"

	[ "${VERBOSITY}" -gt 0 ] && \
		echo "${dev}: pt_start=${pt_start} pt_end=${pt_end}" \
		"pt_size=${pt_size} pt_max=${pt_max} last=${last}"

	# Check if the partition can be grown
	[ "${pt_end}" -eq "${pt_max}" ] && \
		nochange "${dev}: size=${pt_size}, it cannot be grown"
	[ "$((${pt_end} + ${FUDGE}))" -gt "${pt_max}" ] && \
		nochange "${dev}: could only be grown by" \
		"$((${pt_max} - ${pt_end})) [fudge=${FUDGE}]"

	# The partition can be grown if we made it here. Get some more info
	# about it so we can do it properly.
	# FIXME: Do we care about the attribute flags?
	code=$(awk '/^Partition GUID code:/ { print $4 }' "${pt_info}")
	guid=$(awk '/^Partition unique GUID:/ { print $4 }' "${pt_info}")
	name=$(awk '/^Partition name:/ { gsub(/'"'"'/, "") ; \
		if (NF >= 3) print substr($0, index($0, $3)) }' "${pt_info}")
	[ -n "${code}" -a -n "${guid}" ] || \
		fail "${dev}: failed to parse sgdisk details"

	[ "${VERBOSITY}" -gt 0 ] && \
		echo "${dev}: code=${code} guid=${guid} name='${name}'"

	# Calculate the new size of the partition
	new_size=$((${pt_max} - ${pt_start}))
	old="old: size=${pt_size},end=${pt_end}"
	new="new: size=${new_size},end=${pt_max}"
	change_info="${dev}: start=${pt_start} ${old} ${new}"
	
	# Dry run
	[ "${DRY_RUN}" -ne 0 ] && change "${change_info}"

	# Backup the current partition table, we're about to modify it
	sgdisk -b "${GPT_BACKUP}" "${DISK}" > /dev/null || \
		error "${dev}: failed to backup the partition table"

	# Modify the partition table. We do it all in one go (the order is
	# important!):
	#  - move the GPT backup header to the end of the disk
	#  - delete the partition
	#  - recreate the partition with the new size
	#  - set the partition code
	#  - set the partition GUID
	#  - set the partition name
	sgdisk -e -d "${PART}" -n "${PART}":"${pt_start}":"${pt_max}" \
		-t "${PART}":"${code}" -u "${PART}":"${guid}" \
		-c "${PART}":"${name}" "${DISK}" > "${out}" 2>&1 || {
		RESTORE_FUNC=gpt_restore
		cat "${out}" 1>&2
		fail "${dev}: failed to repartition"
	}

	changed "${change_info}"
}

has_cmd() {
	command -v "${1}" >/dev/null 2>&1
}

while [ $# -ne 0 ]; do
	cur=${1}
	next=${2}
	case "$cur" in
		-h|--help)
			Usage
			exit 0
			;;
		--fudge)
			FUDGE=${next}
			shift
			;;
		-N|--dry-run)
			DRY_RUN=1
			;;
		-v|--verbose)
			VERBOSITY=$(($VERBOSITY+1))
			;;
		--)
			shift
			break
			;;
		-*)
			fail "unknown option ${cur}"
			;;
		*)
			if [ -z "${DISK}" ]; then
				DISK=${cur}
			else
				[ -z "${PART}" ] || \
					fail "confused by arg ${cur}"
				PART=${cur}
			fi
			;;
	esac
	shift
done

[ -n "${DISK}" ] || bad_Usage "must supply disk and partition-number"
[ -n "${PART}" ] || bad_Usage "must supply partition-number"

has_cmd "sfdisk" || fail "sfdisk not found"

[ -e "${DISK}" ] || fail "${DISK}: does not exist"

[ "${PART#*[!0-9]}" = "${PART}" ] || fail "partition-number must be a number"

mktemp_d && TEMP_D="${_RET}" || fail "failed to make temp dir"
trap cleanup EXIT

# get the ID of the first partition to determine if it's MBR or GPT
id=$(sfdisk --id --force "${DISK}" 1 2>/dev/null) || \
	fail "unable to determine partition type"

if [ "${id}" = "ee" ] ; then
	has_cmd "sgdisk" || fail "GPT partition found but no sgdisk"
	debug 1 "found GPT partition table (id = ${id})"
	gpt_resize
else
	debug 1 "found MBR partition table (id = ${id})"
	mbr_resize
fi

# vi: ts=4 noexpandtab
