# vim: ft=zsh noet ts=4 sts=4 sw=4

#
# confz functions for dealing with filesystem and mounting
#

# load the zstat builtin and keep stat external
zmodload -F zsh/stat b:zstat

# helper for finding disk drives by their metadata
fs_smartctl_probe() {
	(($+smartctl_probed)) && (($smartctl_probed)) && return 0
	local dev item val rest
	typeset -g smartctl_probed
	typeset -ga smartctl_devices
	typeset -gA smartctl_info
	smartctl --scan | while read dev rest; do
		smartctl_devices+=( $dev )
		smartctl -i $dev | while IFS=':' read item val; do smartctl_info[${dev}:${${item:l}// /_}]=${val/# #}; done
	done
	smartctl_probed=1
}

# helper for probing devices
fs_blkid_probe() {
	unset fs_blkid_output fs_blkid_result
	typeset -gA fs_blkid_output
	typeset -g fs_blkid_result
	local out key val
	fs_blkid_output=( )

	[[ -b $1 ]] || die "no such block device: ${(qqq)1}"

	out=$( blkid -o export -p $1 )
	fs_blkid_result=$?

	case $fs_blkid_result in
		(2);;
		(0) while IFS='=' read key val; do
				[[ $key == *' '* ]] && die "Malformed blkid output for \"export\" format."
				fs_blkid_output[$key]=$val
			done <<<$out;;
		(*) die "blkid probe failed on device: ${(qqq)1}"
	esac
}

# Check that given device is formatted with given type (return 0) or blank (return 1), die otherwise
fs_check_type() {
	local device fstype kind
	device=$1
	fstype=$2
	kind=${3:-TYPE}  # TYPE = filesystem type, PTTYPE = partition format
	fs_blkid_probe $device
	if (($fs_blkid_result)); then
		# empty label
		return 1
	elif (( $+fs_blkid_output[$kind] )); then
			if [[ ${fs_blkid_output[$kind]:-} == $fstype ]]; then
				return 0
			else
				die "$0: non-$fstype label (${fs_blkid_output[$kind]}) already present on device: ${(qqq)device}"
			fi
	else
		local -a labels
		local k v
		for k v in "${(kv)fs_blkid_output[@]}"; do
			# Filter out the information about partition layout rather than content
			case $k in
				(DEVNAME) continue;;
				(PART_ENTRY_*) continue;;
				(*) labels+=( "$k=${(qqq)v}" )
			esac
		done
		if (( $#labels )); then
			die "$0: unexpected filesystem labels present on device: ${(qqq)device} $labels[*]"
		fi
		# empty label
		return 1
	fi
}

# helper for listing partition table
fs_parted_list() {
	unset fs_parted_start fs_parted_end fs_parted_size fs_parted_type
	unset fs_parted_filesystem fs_parted_name fs_parted_flags
	typeset -gA fs_parted_start fs_parted_end fs_parted_size fs_parted_type
	typeset -gA fs_parted_filesystem fs_parted_name fs_parted_flags
	local line got_header number out
	local NumberE StartB StartE EndB EndE SizeB SizeE FSB FSE NameB NameE FlagsB
	local TypeB TypeE
	local MATCH MBEGIN MEND

	[[ -b $1 ]] || die "no such block device: ${(qqq)1}"

	out=$( parted --script $1 -- unit s print ) || \
		die "parted print of device failed: ${(qqq)1}"

	got_header=0
	while IFS= read line; do
		if ((got_header)); then
			number=${${line[1,$NumberE]}// }
			fs_parted_start[$number]=${${line[$StartB,$StartE]}// }
			fs_parted_end[$number]=${${line[$EndB,$EndE]}// }
			fs_parted_size[$number]=${${line[$SizeB,$SizeE]}// }
			fs_parted_flags[$number]=${line[$FlagsB,-1]}

			[[ -n $TypeB ]] && \
				fs_parted_type[$number]=${${line[$TypeB,$TypeE]}// }

			[[ -n $NameB ]] && \
				fs_parted_name[$number]=${${line[$NameB,$NameE]}// }

			[[ -n $FSB ]] && \
				fs_parted_filesystem[$number]=${${line[$FSB,$FSE]}// }

		elif [[ $line == Number* ]]; then
			got_header=1

			[[ $line =~ 'Number +' ]] || die "malformed header: ${(qqq)line}"
			NumberE=$MEND

			[[ $line =~ 'Start +' ]] || die "malformed header: ${(qqq)line}"
			StartB=$MBEGIN
			StartE=$MEND

			[[ $line =~ 'End +' ]] || die "malformed header: ${(qqq)line}"
			EndB=$MBEGIN
			EndE=$MEND

			[[ $line =~ 'Size +' ]] || die "malformed header: ${(qqq)line}"
			SizeB=$MBEGIN
			SizeE=$MEND

			[[ $line =~ 'Flags *' ]] || die "malformed header: ${(qqq)line}"
			FlagsB=$MBEGIN

			if [[ $line =~ 'Type +' ]]; then
				TypeB=$MBEGIN
				TypeE=$MEND
			fi

			if [[ $line =~ 'Name +' ]]; then
				NameB=$MBEGIN
				NameE=$MEND
			fi

			if [[ $line =~ 'File system +' ]]; then
				FSB=$MBEGIN
				FSE=$MEND
			fi
		fi
	done <<<"$out"
}

# Find disks by their IDs
confz_disk_id_check() {
	# arguments: device device_model serial_number
	if ! (($+vars[device] + $+vars[device_model] + $+vars[serial_number])); then
		die "$0: requires one or more arguments: device device_model serial_number"
	fi

	do_command=( true )
	fs_smartctl_probe
	local device
	local -a found

	for device in $smartctl_devices; do
		if (($+vars[device])); then
			[[ $vars[device] != $device ]] && continue
		fi
		if (($+vars[device_model])); then
			[[ $vars[device_model] != $smartctl_info[${device}:device_model] ]] && continue
		fi
		if (($+vars[serial_number])); then
			[[ $vars[serial_number] != $smartctl_info[${device}:serial_number] ]] && continue
		fi
		found+=( $device $smartctl_info[${device}:device_model] $smartctl_info[${device}:serial_number] )
	done

	(( $#found )) || die "$0: found no matching devices"
	(( $#found % 3 )) && die "$0: internal error parsing device information"
	[[ $#found -ne 3 ]] && die "$0: found too many matching devices"
	vars[device]=$found[1]
	vars[device_model]=$found[2]
	vars[serial_number]=$found[3]
}

# check for DOS-style MBR on device
confz_disklabel_dos_check() {
	checkvars device

	do_command=( parted --script $vars[device] -- mklabel msdos )
	fs_check_type $vars[device] dos PTTYPE
}


# check for GPT partition table on device
confz_disklabel_gpt_check() {
	checkvars device

	do_command=( parted --script $vars[device] -- mklabel gpt )
	fs_check_type $vars[device] gpt PTTYPE
}


# embed file in MBR
confz_mbr_code_check() {
	checkvars device from
	do_command=( dd bs=440 count=1 if=$vars[from] of=$vars[device] )
	dd if=$vars[device] bs=440 count=1 | cmp -s - $vars[from]
}


# check for gpt partition
confz_gpt_partition_check() {
	local prev option
	checkvars device number name
	defvar fs ext2
	require disklabel_gpt :device
	fs_parted_list $vars[device]
	(( $+vars[size] || $+vars[end] )) || die "$0: requires either size or end"
	if [[ $vars[number] != 1 ]]; then
		prev=$[ $vars[number] - 1 ]
		(( $+fs_parted_start[$prev] )) || die "$0: missing preceding partition"
		: ${vars[start]:=${fs_parted_start[$prev]}}
	else
		checkvars start
	fi
	if ! (($+vars[end])); then
		vars[end]=$[${${vars[start]}%s} + ${${vars[size]}%s}]s
	fi
	if (( $+fs_parted_start[$vars[number]] )); then
		[[ $fs_parted_name[$vars[number]] == $vars[name] ]] || \
			die "Partition $vars[number] has name ${(qqq)fs_parted_name[$vars[number]}, want ${(qqq)vars[name]}"
		return 0
	else
		do_command=( parted --script )
		(($+vars[align])) && do_command+=( --align $vars[align] )
		do_command+=(
			$vars[device] -- unit s
			mkpart $vars[name] $vars[fs] $vars[start] $vars[end]
		)
		(($+vars[align])) && do_command+=(
			align-check $vars[align] $vars[number]
		)
		(($+vars[options])) && for option in ${=vars[options]}; do
			do_command+=( set $vars[number] $option on )
		done
		return 1
	fi
}


# check for primary partition
confz_primary_partition_check() {
	local prev option
	checkvars device number
	defvar fs ext2
	require disklabel_dos :device
	fs_parted_list $vars[device]
	(( $+vars[size] || $+vars[end] )) || die "$0: requires either size or end"
	if [[ $vars[number] != 1 ]]; then
		prev=$[ $vars[number] - 1 ]
		(( $+fs_parted_start[$prev] )) || die "$0: missing preceding partition"
		: ${vars[start]:=${fs_parted_start[$prev]}}
	else
		checkvars start
	fi
	if ! (($+vars[end])); then
		vars[end]=$[${${vars[start]}%s} + ${${vars[size]}%s}]s
	fi
	if (( $+fs_parted_start[$vars[number]] )); then
		return 0
	else
		do_command=( parted --script )
		(($+vars[align])) && do_command+=( --align $vars[align] )
		do_command+=(
			$vars[device] -- unit s
			mkpart primary $vars[fs] $vars[start] $vars[end]
		)
		(($+vars[align])) && do_command+=(
			align-check $vars[align] $vars[number]
		)
		(($+vars[options])) && for option in ${=vars[options]}; do
			do_command+=( set $vars[number] $option on )
		done
		return 1
	fi
}


# check for bootable flag on primary partition
confz_bootable_partition_check() {
	checkvars device number
	require disklabel_dos :device
	fs_parted_list $vars[device]
	(( $+fs_parted_flags[$vars[number]] )) || \
		die "$0: device ${(q)vars[device]} partition ${(q)vars[number]} not found"
	if [[ $fs_parted_flags[$vars[number]] = *boot* ]]; then
		return 0
	else
		do_command=( parted --script $vars[device] -- set $vars[number] boot on )
		return 1
	fi
}


# create swap partition
confz_swap_check() {
	checkvars device

	do_command=( mkswap $vars[device] )

	fs_check_type $vars[device] swap
}


# enable swap device
confz_swapon_check() {
	local out device rest

	checkvars device
	require swap :device

	out=$(swapon -s) || die "$0: swapon -s comman failed"
	while read device rest; do
		[[ $device == $vars[device] ]] && return 0
	done <<<$out

	do_command=( swapon $vars[device] )
	return 1
}


# set up /dev/md* device
confz_mdraid_check() {
	local line level out seen_level seen_header devices all_empty ret
	local -a seen_devices devices device_numbers
	local -A level_map

	checkvars md_device raid_devices level
	defvar metadata default
	defvar bitmap none

	devices=( "${(Q@)${(z)vars[raid_devices]}}" )

	level_map=(
		linear	linear
		raid0	0
		0	0
		stripe	0
		raid1	1
		1	1
		mirror	1
		raid4	4
		4	4
		raid5	5
		5	5
		raid6	6
		6	6
		raid10	10
		10	10
		multipath	multipath
		mp	multipath
		faulty	faulty
		container	container
	)
	level=${level_map[${vars[level]}]}

	for device in $devices; do
		[[ -b $device ]] || die "$0: not a block device: ${(qqq)device}"
		device_numbers+=( $(zstat +rdev $device) ) || die "$0: could not stat ${(qqq)device}"
	done

	if [[ -b $vars[md_device] ]]; then
		out=$( mdadm --detail $vars[md_device] )
		ret=$?
		case $ret in
			(0) ;;
			(1) ;; # return 1;;
			(*) die "$0: mdadm --detail ${(qqq)vars[md_device]} returned $ret";;
		esac
		seen_header=0
		while read line; do
			if ! (($seen_header)); then
				case $line in
					('Raid Level :'*)
						seen_level=$level_map[${line#* : }];;
					(Number*Major*Minor*RaidDevice*State)
						seen_header=1;;
				esac
			else
				seen_devices+=( $(( ${${=line}[2]} << 8 + ${${=line}[3]} )) )
			fi
		done <<<$out

		if (($seen_header)); then
			[[ $level == $seen_level ]] || \
				die "$0: raid level mismatch." \
				"expected: ${(q)level} got: ${(q)seen_level}"

			[[ ${(j.:.)${(o)device_numbers}} == ${(j.:.)${(o)seen_devices}} ]] || \
				die $0$': raid device mismatch\nexpected:' \
				${(oqqq)device_numbers}$'\ngot:' \
				${(oqqq)seen_devices}

			return 0
		fi
	fi

	all_empty=1
	for device in $devices; do
		# all devices either need to have empty labels or be linux_raid_member
		if fs_check_type $device linux_raid_member; then
			all_empty=0
		fi
	done

	if (($all_empty)); then
		# empty labels
		do_command=(
			mdadm --create
			-l $level
			--metadata=$vars[metadata]
			--bitmap=$vars[bitmap]
			-n $#devices
			$vars[md_device]
			$devices
		)
	else
		# already created
		do_command=( mdadm --assemble $vars[md_device] $devices )
	fi
	return 1
}


# set up LVM2 physical volume
confz_physical_volume_check() {
	checkvars device
	do_command=( lvm pvcreate $vars[device] )
	fs_check_type $vars[device] LVM2_member
}


# configure LVM2 volume group
confz_volume_group_check() {
	local -A devices volumes
	local -a vg_devices
	local device

	[[ -n ${vars[vg_name]:=${DEFAULT_VG}} ]] || \
		die "$0: DEFAULT_VG is unset and no 'vg_name' was passed"

	checkvars vg_devices

	vg_devices=( "${(Q@)${(z)vars[vg_devices]}}" )
	for device in $vg_devices; do
		require physical_volume device=$device
	done

	volumes=( $(lvm vgs --noheadings -o vg_name,pv_count) ) || \
		die "$0: lvm vgs command returned error"
	if ! (($+volumes[${vars[vg_name]}])); then
		do_command=( lvm vgcreate ${vars[vg_name]} $vg_devices )
		return 1
	fi
	devices=( $(lvm pvs --noheadings -o pv_name,vg_name) ) || \
		die "$0: lvm pvs command returned error"
	do_command=( lvm vgextend ${vars[vg_name]} )
	for device in $vg_devices; do
		[[ $+devices[$device] == 1 && $devices[$device] == $vars[vg_name] ]] || \
			do_command+=( $device )
	done
	(( $#do_command == 3 ))
}


# configure LVM2 logical volume
confz_logical_volume_check() {
	[[ -n ${vars[vg_name]:=${DEFAULT_VG}} ]] || \
		die "$0: DEFAULT_VG is unset and no 'vg_name' was passed"

	[[ -n ${vars[size]:=${DEFAULT_VOLUME_SIZE}} ]] || \
		die "$0: DEFAULT_VOLUME_SIZE is unset and no 'size' was passed"

	checkvars lv_name

	setvar device /dev/mapper/$vars[vg_name]-$vars[lv_name]

	do_command=(
		lvcreate
		--name ${vars[lv_name]}
		--size ${vars[size]}
		${vars[vg_name]}
	)

	fail_reason="not a block device: ${vars[device]}"
	[[ -b ${vars[device]} ]]
}

# configure LVM2 logical volume
confz_logical_volume_extents_check() {
	[[ -n ${vars[vg_name]:=${DEFAULT_VG}} ]] || \
		die "$0: DEFAULT_VG is unset and no 'vg_name' was passed"

	checkvars lv_name pv_name extents

	local vg_name lv_name pv_name pvseg_start pvseg_size segtype
	local -A found
	pvs --segments --noheadings --separator $'\t' \
		-o vg_name,lv_name,pv_name,pvseg_start,pvseg_size,segtype \
		| while IFS=$'\t' \
		read vg_name lv_name pv_name pvseg_start pvseg_size segtype
	do
		vg_name=${vg_name/# #}  # strip preceding space
		if matchvars vg_name $vg_name lv_name $lv_name; then
			(($#found)) && die "$0: more than one segment found"
			if [[ $segtype != linear ]] || ! matchvars pv_name $pv_name offset $pvseg_start extents $pvseg_size; then
				die "$0: non-matching physical layout of the volume"
			fi
			found[offset]=$pvseg_start
		fi
	done
	unify "${(kv@)found}" \
		device /dev/mapper/$vars[vg_name]-$vars[lv_name]

	do_command=(
		lvcreate
		--name ${vars[lv_name]}
		--extents ${vars[extents]}
		${vars[vg_name]}
		$vars[pv_name]:$vars[offset]-$[ $vars[offset] + $vars[extents] ]
	)

	if ! (($#found)); then
		fail_reason="not found in pvs output: ${vars[vg_name]}/${vars[lv_name]}"
		return 1
	fi

	fail_reason="not a block device: ${vars[device]}"
	[[ -b ${vars[device]} ]]
}


# create filesystem on block device
confz_filesystem_check() {
	[[ -n ${vars[filesystem]:=${DEFAULT_FS}} ]] || \
		die "$0: DEFAULT_FS is unset and no 'filesystem' was passed"

	checkvars label device filesystem
	defvar mkfs_opts ''

	[[ -b ${vars[device]} ]] || \
		die "$0: not a block device: ${(qqq)vars[device]}"

	do_command=(
		mkfs -t ${vars[filesystem]}
	)
	case $vars[filesystem] in
		(xfs|btrfs|ext[234])
			do_command+=( -L "${vars[label]}" );;
		(reiserfs)
			# Asks questions without -q
			do_command+=( -l "${vars[label]}" -q );;
		(vfat)
			do_command+=( -n "${vars[label]}" );;
		(*)
			if [[ -n $vars[label] ]]; then
				die "$0: I don't know how to set label on ${(qqq)vars[filesystem]}"
			fi
			;;
	esac
	[[ -n $vars[mkfs_opts] ]] && do_command+=( "${(Q@)${(z)vars[mkfs_opts]}}"  )
	do_command+=( ${vars[device]} )

	local tries # blk_out DEVNAME LABEL UUID TYPE SEC_TYPE PARTLABEL PARTUUID

	tries=10
	while ((tries)); do
		fs_blkid_probe $vars[device]
		[[ $fs_blkid_result == 0 ]] && break
		tries=$[$tries - 1]
	done

	fail_reason="no blkid signature found on $vars[device]"
	# Check if anything was detected at all
	(( $#fs_blkid_output )) || return 1

	# Check if it is a filesystem
	[[ -z ${fs_blkid_output[TYPE]:-} && -z ${fs_blkid_output[LABEL]:-} && -z ${fs_blkid_output[UUID]:-} ]] && return 1

	# Check if it matches requirements
	[[ ${fs_blkid_output[LABEL]:-} == ${vars[label]:-} && $fs_blkid_output[TYPE] == $vars[filesystem] ]] && return 0

	die "$0: filesystem already present on ${(qqq)vars[device]}!"
}


# put mountpoint for device into /etc/fstab
confz_fstab_check() {
	checkvars device mountpoint filesystem
	defvar opts noatime
	defvar dump 0
	defvar fstab /etc/fstab
	if [[ $vars[filesystem] == xfs ]]; then
		defvar pass 0
	else
		defvar pass 2
	fi

	local device mountpoint filesystem opts dump pass
	sed '/^[ \t]*#/d;s/#.*//;s/[ \t]\+/ /g' $vars[fstab] | \
	while read device mountpoint filesystem opts dump pass; do
		if [[ $mountpoint == ${vars[mountpoint]} ]]; then
			[[ $device == ${vars[device]} && opts != *bind* ]] || \
				die "$0: $mountpoint already present" \
					"with different device than ${(qqq)vars[device]}"

			[[ $filesystem == ${vars[filesystem]} ]] || \
				die "$0: $mountpoint already present" \
					"with different filesystem than ${(qqq)vars[filesystem]}"

			[[ $opts == ${vars[opts]} ]] || \
				die "$0: $mountpoint already present" \
					"with different opts than ${(qqq)vars[opts]}"

			[[ $dump == ${vars[dump]} ]] || \
				die "$0: $mountpoint already present" \
					"with different dump than ${(qqq)vars[dump]}"

			[[ $pass == ${vars[pass]} ]] || \
				die "$0: $mountpoint already present" \
					"with different pass than ${(qqq)vars[pass]}"

			return 0  # found matching entry
		fi
	done

	fail_reason="no entry for ${(qqq)vars[mountpoint]} in ${(qqq)vars[fstab]}"
	return 1  # did not find matching entry
}

confz_fstab_do() {
	print -r - >>$vars[fstab] ${vars[device]}$'\t'${vars[mountpoint]}$'\t'${vars[filesystem]}$'\t'${vars[opts]}$'\t'${vars[dump]}' '${vars[pass]}
}


# make device mounted on mountpoint
confz_mounted_check() {
	checkvars device mountpoint

	fail_reason="could not find ${vars[device]} ${vars[mountpoint]} in /proc/mounts"
	grep -q "^${vars[device]} ${vars[mountpoint]} " /proc/mounts
}

confz_mounted_do() {
	mkdir -p ${vars[mountpoint]} || return $?
	mount "$@" ${vars[device]} ${vars[mountpoint]}
}


# make device mounted on mountpoint
confz_bind_mounted_check() {
	checkvars device origin mountpoint

	fail_reason="could not find ${vars[device]} ${vars[mountpoint]} in /proc/mounts"
	grep -q "^${vars[device]} ${vars[mountpoint]} " /proc/mounts
}

confz_bind_mounted_do() {
	mkdir -p ${vars[mountpoint]} || return $?
	mount --bind "$@" ${vars[origin]} ${vars[mountpoint]}
}


# create LVM2 logical volume, and make sure it's in fstab and mounted
confz_mounted_volume_check() {
	checkvars mountpoint
	defvar lv_name ${${${vars[mountpoint]}##/}//\//_}
	defvar filesystem xfs
	defvar label ${${vars[lv_name]}[1,12]}
	defvar root /

	[[ $vars[mountpoint] == /* ]] || \
			die "$0: mountpoint must be absolute path, got: ${(qqq)vars[mountpoint]}"
	require logical_volume %device \?vg_name :size :lv_name
	require filesystem :device :label :filesystem \?mkfs_opts
	require fstab fstab=${vars[root]%/}/etc/fstab \
		:device :mountpoint :filesystem \?opts \?dump \?pass
	require mounted :device mountpoint=${vars[root]%/}/${vars[mountpoint]#/}

	do_command=( true )
}


confz_syslinux_modules_check() {
	checkvars directory
	if ! [[ -e $vars[directory]/menu.c32 ]]; then
		do_command=( cp -v /usr/share/syslinux/*.c32 $vars[directory]/ )
		return 1
	else
		return 0
	fi
}

# install extlinux
confz_extlinux_check() {
	checkvars directory
	defvar extlinux /sbin/extlinux
	defvar install_touch_file $vars[directory]/.extlinux_installed
	[[ -e $vars[directory]/extlinux.conf || -e $vars[directory]/syslinux.cfg ]] || \
		die "no configuration file for extlinux found"
	require syslinux_modules :directory
	[[ -e $vars[install_touch_file] ]]
	return $?
}

confz_extlinux_do() {
	$vars[extlinux] -i -r $vars[directory] || return $?
	touch $vars[install_touch_file] || return $?
}


# install GRUB2
confz_grub2_check() {
	checkvars device
	defvar boot_directory /boot
	defvar install_touch_file vars[boot_directory]/.grub2${${vars[device]}//\//.}
	[[ -e $vars[install_touch_file] ]]
	return $?
}

confz_grub2_do() {
	grub2-install --boot-directory=$vars[boot_directory] $vars[device] || return $?
	touch $vars[install_touch_file] || return $?
}