lcss

#!/bin/bash

if [ "$USER" != "root" ]; then
	echo "$USER is not root." >&2
	exit 1
fi

if [ -z "$1" ]; then
	echo "$0 IMAGE.img [SCRIPT] [EXPAND] [SHRINK]" >&2
	exit 1
fi

if [ ! -f "$1" ]; then
	echo "$1 is not a file." >&2
	exit 1
fi

if [ ! -z "$2" ]; then
	if [ ! -f "$2" ]; then
		echo "$2 is not a file." >&2
		exit 1
	fi
	script_bin=$(head -n 1 "$2")
	if [ "${#script_bin}" -gt 100 ]; then
		echo "$2 is not an interpreted script." >&2
		exit 1
	elif [ "${script_bin:0:2}" != "#!" ]; then
		echo "$2 is missing the interpreter header." >&2
		exit 1
	fi
	script_int_path="${script_bin:2}"
	script="$(readlink -f "$2")"
fi

if [ -z "$3" ]; then
	image_expand=0
else
	image_expand="${3^^}"
	if ! echo "$image_expand" | grep -E "^[[:digit:]]+[MGT]?$" > /dev/null; then
		echo "EXPAND must be an integer optionally followed by [MGT] unit, M assumed." >&2
		exit 1
	fi
	image_expand_num=$(echo "$image_expand" | grep -oE "^[[:digit:]]+")
	image_expand_exp=$(echo "$image_expand" | grep -oE "[MGT]?$")
	if [ -z "$image_expand_exp" ]; then
		image_expand=${image_expand}M
	fi
	#if [ -z "$image_expand_exp" ]; then
	#	if [ "$image_expand_num" -lt $((1<<20)) ]; then
	#		echo "EXPAND must be greater than 1M." >&2
	#		exit 1
	#	fi
	#	if [ "$image_expand_num" -ne "$((image_expand_num>>20<<20))" ]; then
	#		echo "EXPAND must be a multiple of 1M." >&2
	#		exit 1
	#	fi
	#fi
	image_expand="+$image_expand"
fi

if [ -z "$4" ]; then
	image_shrink=0
else
	image_shrink="${4^^}"
	if ! echo "$image_shrink" | grep -E "^[[:digit:]]+[MGT]?$" > /dev/null; then
		echo "SHRINK must be an integer optionally followed by [MGT] unit, M assumed." >&2
		exit 1
	fi
	image_shrink_num=$(echo "$image_shrink" | grep -oE "^[[:digit:]]+")
	image_shrink_exp=$(echo "$image_shrink" | grep -oE "[MGT]?$")
	if [ -z "$image_shrink_exp" ]; then
		if [ "$image_shrink_num" -lt $((1<<6)) ]; then
			echo "SHRINK must be greater than 64M." >&2
			exit 1
		fi
	elif [ "$image_shrink_exp" = "M" ]; then
		if [ "$image_shrink_num" -lt $((1<<6)) ]; then
			echo "SHRINK must be greater than 64M." >&2
			exit 1
		fi
		$image_shrink=$image_shrink_num
	elif [ "$image_shrink_exp" = "G" ]; then
		$image_shrink=$((image_shrink_num<<10))
	elif [ "$image_shrink_exp" = "T" ]; then
		$image_shrink=$((image_shrink_num<<20))
	else
		echo "SHRINK unexpected assert error for value $image_shrink." >&2
		exit 1
	fi
fi

image_file="$(readlink -f "$1")"

cd $(readlink -f $(dirname ${BASH_SOURCE[0]}))

. lib/chroot.sh
. lib/loop.sh
. lib/traps.sh

set -e

PT_TYPE_MBR='mbr'
PT_TYPE_GPT='gpt'
PART_TYPE_EFI='07'
PART_TYPE_LINUX='83'
PART_TYPE_MBR_EFI_GREP=', type=\(ef\|07\|0b\|0c\)'
PART_TYPE_MBR_LINUX_GREP=', type=83'
PART_TYPE_GPT_EFI_GREP='\(EF00 EFI\|0700 Microsoft\)'
PART_TYPE_GPT_LINUX_GREP='8300 Linux'

image_pt_type_str=$(file -bks "$image_file")

if [ "${image_pt_type_str/GPT partition table/}" != "$image_pt_type_str" ]; then
	image_pt_type=$PT_TYPE_GPT
	image_pt_tool="sgdisk -p"
elif [ "${image_pt_type_str/MBR boot sector/}" != "$image_pt_type" ]; then
	image_pt_type=$PT_TYPE_MBR
	image_pt_tool="sfdisk -d"
else
	echo "Partition type not supported." >&2
	exit 1
fi

if ! which ${image_pt_tool%% *} > /dev/null; then
	echo "$image_pt_tool is not available." >&2
	exit 1
fi

image_pt_table=$($image_pt_tool "$image_file")

PT_getPartNum(){
	case $image_pt_type in
		$PT_TYPE_MBR)
			if [ $1 = $PART_TYPE_EFI ]; then
				local part_type_grep="$PART_TYPE_MBR_EFI_GREP"
			elif [ $1 = $PART_TYPE_LINUX ]; then
				local part_type_grep="$PART_TYPE_MBR_LINUX_GREP"
			else
				echo "Unsupported partition type $1." >&2
				return 1
			fi
			local part_num=$(echo "$image_pt_table" | tr -s ' ' | grep "$part_type_grep" | sed "s/^${image_file//\//\\\/}//" | cut -f 1 -d ' ')
			if [ -z "$part_num" ]; then
				echo "Unable to find partition number." >&2
				return 1
			fi
			echo -n "$part_num"
			;;
		$PT_TYPE_GPT)
			if [ $1 = $PART_TYPE_EFI ]; then
				local part_type_grep="$PART_TYPE_GPT_EFI_GREP"
			elif [ $1 = $PART_TYPE_LINUX ]; then
				local part_type_grep="$PART_TYPE_GPT_LINUX_GREP"
			else
				echo "Unsupported partition type $1." >&2
				return 1
			fi
			local part_offset=$(echo "$image_pt_table" | tr -s ' ' | grep "$part_type_grep" | (echo -n ' ' && cat) | tr -s ' ' | cut -f 3 -d ' ')
			if [ -z "$part_num" ]; then
				echo "Unable to find partition number." >&2
				return 1
			fi
			echo -n "$part_offset"
			;;
		*)
			echo "Unsupported partition table type $image_pt_type." >&2
			return 1
			;;
	esac
}

image_root_num=$(PT_getPartNum $PART_TYPE_LINUX)
image_boot_num=$(PT_getPartNum $PART_TYPE_EFI)

LMOUNT_main(){
	traps_start
	
	if [ "$image_expand" != "0" ]; then
		truncate -s "$image_expand" "$image_file"
	fi
	

	local loop_dev=$(LOOP_setup "$image_file")
	traps_push LOOP_detach "$loop_dev"
	
	local loop_traps_length=$TRAPS_LENGTH
	
	if [ ! -b "${loop_dev}p${image_root_num}" ]; then
		partprobe "$loop_dev"
	fi
	
	if [ "$image_expand" != "0" ]; then
		echo ", +" | ${image_pt_tool%% *} -N "$image_root_num" "$loop_dev"
		partprobe "$loop_dev"
	fi

	if [ ! -b "${loop_dev}p${image_root_num}" ]; then
		echo "Loop device partition $image_root_num did not enumerate." >&2
		exit 1
	fi
	
	if [ ! -b "${loop_dev}p${image_root_num}" ]; then
		echo "Loop device partition $image_boot_num did not enumerate." >&2
		exit 1
	fi

	local image_mount_dir=$(mktemp -d)
	traps_push rmdir "$image_mount_dir"
	mount -o noatime,compress=zstd,nossd_spread "${loop_dev}p${image_root_num}" "$image_mount_dir"
	traps_push sleep 1
	traps_push umount -l "$image_mount_dir"

	if [ ! -d "$image_mount_dir/boot/efi" ]; then
		echo "There is no /boot/efi directory in the Linux root." >&2
		exit 1
	fi

	#mount "${loop_dev}p${image_boot_num}" "$image_mount_dir/boot/efi"
	#traps_push sleep 1
	#traps_push umount -l "$image_mount_dir/boot/efi"

	CHROOT_mount "$image_mount_dir"
	
	if [ "$image_expand" != "0" ]; then
		btrfs filesystem resize max "$image_mount_dir"
	fi

	chroot "$image_mount_dir" mount "${loop_dev}p${image_boot_num}" /boot/efi
	traps_push chroot "$image_mount_dir" umount /boot/efi
	
	if [ ! -z "$script" ]; then
		if [ ! -e "${image_mount_dir}${script_int_path}" ]; then
			echo "The script interpreter does not exist on the image." >&2
			exit 1
		elif [ ! -x "${image_mount_dir}${script_int_path}" ]; then
			echo "The script interpreter has no executive privileges." >&2
			exit 1
		fi
		chroot "$image_mount_dir" "$script_int_path" < "$script"
	else
		chroot "$image_mount_dir"
	fi
	
	if [ "$image_shrink" != "0" ]; then
		#local image_linux_size_b=$(btrfs filesystem show --raw "$image_mount_dir" | grep -Eo "size\\s+[0-9]+" | cut -d " " -f 2)
		#local image_linux_size=$((image_linux_size_b >> 20))
		#local image_linux_sect=$((image_linux_size_b >> 9))
		local image_linux_shrink_step=64
		local image_shrunk=0
		while true; do
			local image_linux_unalloc_b=$(btrfs filesystem usage -b "$image_mount_dir" | grep "Device unallocated" | grep -Eo "[0-9]*")
			local image_linux_unalloc=$((image_linux_unalloc_b >> 20))
			if [ "$((image_linux_unalloc - $image_linux_shrink_step))" -le "$image_shrink" ]; then
				break
			fi
			if ! btrfs filesystem resize "-${image_linux_shrink_step}M" "$image_mount_dir"; then
				break
			fi
			local image_shrunk=$((image_shrunk+$image_linux_shrink_step))
		done
		
		traps_popUntilLength "$loop_traps_length"
		
		echo ", -${image_shrunk}M" | ${image_pt_tool%% *} -N "$image_root_num" "$loop_dev"
	
	fi

	traps_exit || true
	if [ "$image_shrink" != "0" ]; then
		truncate -s "-${image_shrunk}M" "$image_file"
		# move GPT backup table
	fi
}

LMOUNT_main