#!/usr/bin/env bash
_scriptname="img-shrink"
set -u
set -e

function println { printf '%s\n' "$*"; }
function error { println "Error: $*" >&2; }
function step { println "==> $*"; }


################################################################################
# Parse input                                                                  #
################################################################################
image="${1:-}"
if [[ -z "$image" ]]; then
	println "Usage: $_scriptname <disk_image>"
	exit 0
fi

if ! [[ -f "$image" ]]; then
	error "Not a regular file: $image"
	exit 1
fi


################################################################################
# Create a loop device and map the partitions                                  #
################################################################################
function map_image
{
	step "Mapping image"
	loop=$(losetup -fv --show "$image")
	if ! [[ "$loop" =~ ^(/dev/[^[:blank:]]+)$ ]]; then
		error "Failed to set up loop device"
		exit 1
	fi
	loop="${BASH_REMATCH[1]}"
	println "Loop device: $loop"
	partitions=()
	while read -r p; do
		if [[ $p =~ [[:blank:]](loop[[:alnum:]]+) ]]; then
			partitions+="${BASH_REMATCH[1]}"
		fi
	done < <(kpartx -avs "$loop")
	if ! (( ${#partitions[@]} )); then
		error "No partitions found"
		exit 1
	fi
	println "Partitions: ${partitions[*]}"
	for ((i=0; i<${#partitions[@]}; i++)); do
		partitions[$i]="/dev/mapper/${partitions[$i]}"
	done
	last="${partitions[-1]}"
}
map_image


################################################################################
# Make sure we clean up the trash when we're done                              #
################################################################################
tmp=$(mktemp -d)
function unmap_image
{
	for p in "${partitions[@]}"; do
		umount -f "$p" || :
	done
	sync
	local i
	for ((i=0; i<10; i++)); do
		if losetup -a | grep "^${loop}:"; then
			kpartx -dv "$loop" >/dev/null 2>&1 || :
			losetup -d "$loop" >/dev/null 2>&1 || :
		else
			break
		fi
	done
	if losetup -a | grep "^${loop}:"; then
		error "Failed to unbind ${loop}"
		exit 1
	fi
}
function cleanup
{
	trap : INT TERM
	step "Unmapping image"
	unmap_image
	rmdir "$tmp"
}
trap cleanup EXIT


################################################################################
# Print partition types, mostly for debugging                                  #
################################################################################
step "Testing partitions"
for p in "${partitions[@]}"; do
	mount "$p" "$tmp"
	println "${p}: $(findmnt "$tmp" -o fstype -n)"
	umount "$p"
done


################################################################################
# Inspecting the last partition                                                #
################################################################################
mount "$last" "$tmp"
fstype="$(findmnt "$tmp" -o fstype -n)"
case "$fstype" in
	"ext3"|"ext4") ;;
	*)
		error "Unsupported file system type for last partition: $fstype"
		exit 1
	;;
esac
umount "$last"


################################################################################
# Inspect all partitions                                                       #
################################################################################
step "Checking"
for p in "${partitions[@]}"; do
	e2fsck -f -y "$p"
done


################################################################################
# Defragment for better efficiency                                             #
################################################################################
step "Defragmenting"
for p in "${partitions[@]}"; do
	mount "$p" "$tmp"
	case "$(findmnt "$tmp" -o fstype -n)" in
		"ext3"|"ext4")
			e4defrag "$p" >/dev/null
		;;
	esac
	umount "$p"
done


################################################################################
# Inspect all partitions (again)                                               #
################################################################################
step "Checking"
for p in "${partitions[@]}"; do
	e2fsck -f -y "$p"
done


################################################################################
# Shrink the last file system                                                  #
################################################################################
step "Shrinking file system on last partition"
resize2fs -M "${partitions[-1]}"


################################################################################
# Inspect all partitions (again)                                               #
################################################################################
step "Checking"
for p in "${partitions[@]}"; do
	e2fsck -f -y "$p"
done


################################################################################
# Shrink the last partition                                                    #
################################################################################
step "Resize root partition"
root_count=""
root_size=""
root_begin=""
root_sector=""
while read -r line; do
	case "${line,,}" in
		"block count: "*) root_count=${line##* } ;;
		"block size: "*) root_size=${line##* } ;;
	esac
done < <(dumpe2fs -h "$last")
while read -r line; do
	case "${line,,}" in
		*'/'"${last##*/} "*)
			if [[ "$line" =~ \* ]]; then
				# boot drive
				root_begin=$(awk '{print $3}' <<<"$line")
				root_end_old=$(awk '{print $4}' <<<"$line")
			else
				root_begin=$(awk '{print $2}' <<<"$line")
				root_end_old=$(awk '{print $3}' <<<"$line")
			fi
		;;
		"units: sectors of "*) root_sector=$(awk '{print $1}' <<< "${line##* = }") ;;
	esac
done < <(fdisk -l "$loop")
if [[ -z "$root_count" ]] || [[ -z "$root_size" ]] || [[ -z "$root_begin" ]] || [[ -z "$root_sector" ]]; then
	echo "Unable to determine root file system size :(" >&2
	exit 1
fi
root_end=$((root_begin + (root_count*root_size/root_sector) - 1))
echo "Root FS block count:             $root_count"
echo "Root FS block size:              $root_size"
echo "Root FS begins at:               $root_begin"
echo "Root FS sector size:             $root_sector"
echo "Root FS partition ends at:       $root_end_old"
echo "Root FS partition should end at: $root_end"
if [[ "$root_end_old" = "$root_end" ]]; then
	println "Partition size is already correct"
else
	parted "${loop}" unit s resizepart "${last##*p}" "${root_end}s" Yes
fi


################################################################################
# Zero fill the free space                                                     #
################################################################################
step "Zero filling"
for p in "${partitions[@]}"; do
	println "$p"
	mount "$p" "$tmp"
	cat /dev/zero >"${tmp}/.zerofill" 2>/dev/null || :
	rm "${tmp}/.zerofill"
	umount "$p"
done


################################################################################
# Truncate the image                                                           #
################################################################################
step "Truncating"
unmap_image
truncate -s "$(( ( root_end + 1 ) * root_sector ))" "$image"


################################################################################
# Final inspection                                                             #
################################################################################
map_image
step "Final inspection"
for p in "${partitions[@]}"; do
	e2fsck -f -y "$p"
done



# vim: tabstop=4:softtabstop=4:shiftwidth=4:noexpandtab
