#! /bin/sh -e

. /usr/share/debconf/confmodule

# Be nice to monolithic users.
db_get rescue/enable
[ "$RET" = true ] || exit 0

log () {
	logger -t rescue-mode "$@"
}

try_load_module () {
	log-output -t rescue modprobe "$1" || true
}

scan_lvm () {
	log-output -t rescue pvscan || true
	log-output -t rescue vgscan || true
}

get_passphrase () {
	db_set rescue/passphrase ""
	db_fset rescue/passphrase seen false
	db_subst rescue/passphrase DEVICE "$dev"
	db_input critical rescue/passphrase

	db_go || return 1

	db_get rescue/passphrase || RET=''
	echo -n "$RET"
}

do_cryptsetup () {
	local pass_ok pass dev cryptdev

	dev="$1"
	pass_ok=0
	cryptdev="${dev##*/}_crypt"
	if ! cryptsetup status "$cryptdev" > /dev/null; then
		while [ $pass_ok -eq 0 ]; do
			pass="$(get_passphrase)" || return 1
			if [ -z "$pass" ]; then
				return 1
			fi
			echo -n "$pass" | log-output -t rescue \
				cryptsetup -d - luksOpen "$dev" "$cryptdev" &&
				pass_ok=1
		done
	fi
	CRYPTPARTS="${CRYPTPARTS:+$CRYPTPARTS, }/dev/mapper/$cryptdev"
}

activate_encrypted_partitions () {
	for dev in $(list-devices partition); do
		if cryptsetup isLuks "$dev" 2> /dev/null; then
			do_cryptsetup "$dev"
		fi
	done
	
	if [ -z "$CRYPTPARTS" ]; then
		return 1
	fi
	try_load_module aes
}

try_load_module btrfs
try_load_module ext2
try_load_module ext3
try_load_module ext4
try_load_module jfs
try_load_module xfs

# Linux root filesystems won't be on vfat, but this may be useful anyway ...
try_load_module vfat

# RAID support
try_load_module md
try_load_module raid0
try_load_module raid1
try_load_module raid5   # <=2.6.17
try_load_module raid456 # >=2.6.18
# mdadm will fail if /dev/md does not already exist
mkdir -p /dev/md
# Don't auto-assemble; there'll be a menu entry for this.

# LVM support
try_load_module dm-mod
try_load_module lvm-mod
scan_lvm

# Crypto support
CRYPTPARTS=
if activate_encrypted_partitions; then
	# Scan for LVM partitions again as we just added encrypted volumes
	scan_lvm
fi

db_capb backup

TARGET_MOUNTS=

# Try to unmount everything we previously mounted.
umount_all () {
	for old_mount in $TARGET_MOUNTS; do
		log-output -t rescue umount /target"$old_mount" || true
	done
	TARGET_MOUNTS=
	log-output -t rescue umount /target || true
}

trap umount_all EXIT HUP INT QUIT TERM

# Given a path, bind-mounts it into /target
mount_bind_fs () {
	local mountpoint

	mountpoint=$1
	if [ -d /target"$mountpoint" ]; then
		if log-output -t rescue mount -o bind "$mountpoint" /target"$mountpoint"; then
			TARGET_MOUNTS="${mountpoint}${TARGET_MOUNTS:+ $TARGET_MOUNTS}"
		fi
	fi
}

# Choose a root filesystem.
choose_root () {
	PARTITIONS="$(list-devices partition | sort | \
		tr '\n' , | sed 's/,$//; s/,/, /g')"

	if [ "$CRYPTPARTS" ]; then
		PARTITIONS="${PARTITIONS:+$PARTITIONS, }$CRYPTPARTS"
	fi

	if lvdisplay -c | cut -d: -f4 | grep 0 >/dev/null; then
		# Some unavailable logical volumes; activate them
		log-output -t rescue vgchange -a y || true
	fi
	LVMPARTS="$(lvdisplay -c | sed 's/^ *//' | cut -d: -f1 | sort | \
		tr '\n' , | sed 's/,$//; s/,/, /g')" || LVMPARTS=
	if [ "$LVMPARTS" ]; then
		PARTITIONS="${PARTITIONS:+$PARTITIONS, }$LVMPARTS"
	fi

	for i in $(grep ^md /proc/mdstat | grep -v inactive | \
		   sed -e 's/^\(md.*\) : active \([[:alnum:]]*\).*/\1/'); do
		NUMBER="${i#md}"
		if [ -e "/dev/md/$NUMBER" ]; then
			MDPART="/dev/md/$NUMBER"
		else
			MDPART="/dev/md$NUMBER"
		fi
		PARTITIONS="${PARTITIONS:+$PARTITIONS, }$MDPART"
	done

	# Deduplicate and sort entries
	PARTITIONS="$(echo "$PARTITIONS" | sed 's/, /\n/g' | sort -u | \
		tr '\n' ',' | sed 's/,$//; s/,/, /g')"

	if [ -z "$PARTITIONS" ]; then
		log "no partitions found!"
		db_input critical rescue/no-partitions
		db_go || exit 10
		export RESCUE_ROOTDEV=
		return
	fi
	log "partitions found: $PARTITIONS"
	db_subst rescue/root PARTITIONS "$PARTITIONS"

	db_input critical rescue/root
	db_go || exit 10
	db_get rescue/root
	export RESCUE_ROOTDEV="$RET"
	log "selected root device '$RESCUE_ROOTDEV'"
	if [ "$RESCUE_ROOTDEV" = none ] || \
	   [ "$RESCUE_ROOTDEV" = assemble-raid ]; then
		umount_all
		return
	fi
	if [ ! -e "$RESCUE_ROOTDEV" ]; then
		log "'$RESCUE_ROOTDEV' does not exist"
		db_subst rescue/no-such-device DEVICE "$RESCUE_ROOTDEV"
		db_input critical rescue/no-such-device
		# Since continuing returns to the rescue/root question,
		# backing up returns to the main menu.
		db_go || exit 10
		return
	fi

	mkdir -p /target
	umount_all
	if ! log-output -t rescue mount "$RESCUE_ROOTDEV" /target; then
		log "mount '$RESCUE_ROOTDEV' /target failed"
		db_subst rescue/mount-failed DEVICE "$RESCUE_ROOTDEV"
		db_input critical rescue/mount-failed
		# Since continuing returns to the rescue/root question,
		# backing up returns to the main menu.
		db_go || exit 10
		return
	fi

	# We want this kernel's idea of device numbering, not
	# whatever might happen to be statically configured on this
	# filesystem.
	mount_bind_fs /dev

	# Virtual filesystems.
	mount_bind_fs /proc
	mount_bind_fs /sys
	mount_bind_fs /run
}

# decode 3-digit octal sequences \nnn in fstab field
decode_octal () {
	printf "$(printf '%s' "$1" | sed -E 's/%|\\/&&/g; s/\\([0-7]{3})/\1/g')"
}

mount_separate_fs () {
	local filesystem fs_dev
	local dev mtpt fstype options rest

	fs_dev=''
	filesystem="$1"
	if [ -f /target/etc/fstab ]; then
		while read -r dev mtpt fstype options rest; do
			if [ -n "${dev##\#*}" ] && [ "$(decode_octal "$mtpt")" = "$filesystem" ]; then
				logger -t rescue "found dev $dev for mtpt $mtpt"
				fs_dev=$(decode_octal "$dev")
				break
			fi
		done </target/etc/fstab
	fi

	[ -n "$fs_dev" ] || return 0
	# If given NAME=<foo>, we need to convert here
	case $fs_dev in
		UUID=*|LABEL=*|PARTUUID=*|PARTLABEL=*)
			fs_dev=$(blkid -t "$fs_dev" -l -o device)
			;;
	esac
	db_subst rescue/separate-fs FILESYSTEM "${filesystem}"
	db_input critical rescue/separate-fs
	db_go || return 1
	db_get rescue/separate-fs
	if [ "$RET" = true ]; then
		if log-output -t rescue mount -w ${options:+-o} "$options" "${fs_dev}" "/target${filesystem}"; then
			# prepend ${filesystem}, to provide reverse ordering to umount_all
			TARGET_MOUNTS="${filesystem}${TARGET_MOUNTS:+ $TARGET_MOUNTS}"
		fi
	fi
	return 0
}

# Prepare a menu of the things you can do to this root filesystem.
prep_menu () {
	mkdir -p /var/lib/rescue
	>/var/lib/rescue/menu-mapping
	CHOICES=
	ITEMS="$(find /lib/rescue.d/ -type f -perm -100 | sort)"
	for item in $ITEMS; do
		base="${item##*/}"
		name="${base#[0-9][0-9]}"
		case $name in
			*.*)	continue ;;
		esac
		if [ -x "$item.tst" ] && ! "$item.tst"; then
			continue
		fi
		db_subst "rescue/menu/$name" DEVICE "$RESCUE_ROOTDEV" \
			|| continue
		db_metaget "rescue/menu/$name" description || continue
		echo "$base:$RET" >> /var/lib/rescue/menu-mapping
		RET="$(echo "$RET" | sed 's/,/\\,/g')"
		CHOICES="${CHOICES:+$CHOICES, }$RET"
	done
	db_subst rescue/menu CHOICES "$CHOICES"
	db_subst rescue/menu DEVICE "$RESCUE_ROOTDEV"
}

# Work out which menu item was selected.
get_menu_item () {
	newline='
'
	IFS_SAVE="$IFS"
	IFS="$newline"
	for line in $(cat /var/lib/rescue/menu-mapping); do
		IFS="$IFS_SAVE"
		if [ "${line#*:}" = "$1" ]; then
			echo "${line%%:*}"
			return
		fi
		IFS="$newline"
	done
	IFS="$IFS_SAVE"
}

# Assemble a RAID array.
assemble_raid () {
	PARTITIONS="$(list-devices partition | sort | \
		tr '\n' , | sed 's/,$//; s/,/, /g')"

	if [ "$CRYPTPARTS" ]; then
		PARTITIONS="${PARTITIONS:+$PARTITIONS, }$CRYPTPARTS"
	fi

	if [ -z "$PARTITIONS" ]; then
		log "no partitions found!"
		db_input critical rescue/no-partitions
		db_go || exit 10
		export RESCUE_ROOTDEV=
		return
	fi
	log "assemble_raid: partitions found: $PARTITIONS"
	db_subst rescue/assemble-raid PARTITIONS "$PARTITIONS"
	db_set rescue/assemble-raid ''

	db_input critical rescue/assemble-raid
	db_go || return 0
	db_get rescue/assemble-raid
	MDADM_CONFIG=
	for raid_partition in $(echo "$RET" | sed 's/[, ][, ]*/ /g'); do
		case $raid_partition in
			auto)
				MDADM_CONFIG='--config=partitions'
				break
				;;
			*)
				MDADM_CONFIG="${MDADM_CONFIG:+$MDADM_CONFIG }$raid_partition"
				;;
		esac
	done
	log-output -t rescue --pass-stdout \
		   mdadm --examine --scan "$MDADM_CONFIG" > /tmp/mdadm.conf || true
	log-output -t rescue \
		   mdadm --assemble --scan --run --config=/tmp/mdadm.conf --auto=yes || true
	scan_lvm
}

STATE=1
while :; do
	case $STATE in
		1)
			choose_root
			if [ "$RESCUE_ROOTDEV" = assemble-raid ]; then
				STATE=4
			else
				STATE=2
			fi
			;;
		2)
			if mount_separate_fs /usr && \
			   mount_separate_fs /boot && \
			   mount_separate_fs /boot/efi ; then
				STATE=3
			else
				STATE=1
			fi
			;;
		3)
			prep_menu
			db_input critical rescue/menu
			if ! db_go; then
				STATE=1
				continue
			fi
			db_get rescue/menu
			item="$(get_menu_item "$RET")"
			if [ -z "$item" ]; then
				log "Could not find menu item for '$RET'!"
				continue
			fi
			item_name="${item#[0-9][0-9]}"
			if ([ -z "$RESCUE_ROOTDEV" ] || \
			    [ "$RESCUE_ROOTDEV" = none ]) && \
			   db_get "rescue/$item_name/intro/no-target"; then
				db_input high "rescue/$item_name/intro/no-target" || true
				if ! db_go; then
					continue
				fi
			elif db_subst "rescue/$item_name/intro" \
				      DEVICE "${RESCUE_ROOTDEV}"; then
				db_input high "rescue/$item_name/intro" || true
				if ! db_go; then
					continue
				fi
			fi
			code=0
			"/lib/rescue.d/$item" || code=$?
			case $code in
				0)	: ;;
				10)	STATE=1 ;;
				*)
					# This is just a fallback; on error,
					# rescue operations should display a
					# more specific error message and
					# exit zero to indicate that the
					# error has been handled.
					db_capb
					db_subst rescue/menu-error OPERATION "${item#[0-9][0-9]}"
					db_subst rescue/menu-error CODE "$code"
					db_input critical rescue/menu-error
					db_go || true
					db_capb backup
					;;
			esac
			;;
		4)
			assemble_raid
			STATE=1
			;;
		*)
			exit 10
			;;
	esac
done