#!/bin/bash
# script can be started by links /usr/sbin/luksunlock or /usr/sbin/lock
# ACTION is a link name (like busybox or kmod)
ACTION=$(basename $0)
LOG=$(mktemp -p /tmp ${ACTION}.XXXX)
EXITCODE=0
REWRITE="\e[25D\e[1A\e[K"

export TEXTDOMAINDIR=/usr/share/locale
export TEXTDOMAIN=luksunlock

echo_() {
	gettext -s "$@"
}

getpass() {
	echo_ "Please enter password for this LUKS partition:" 1>&2
	while IFS= read -r -s -n1 char; do
		if [[ -z "$char" ]]; then
			echo 1>&2
			break
		fi
		if [[ "$char" == $'\x7f' ]]; then
			if [ -n "$PASS" ]; then
			PASS="${PASS%?}"
			echo -n $'\b \b' 1>&2
			fi
		else
			PASS+="$char"
			echo -n '*' 1>&2
		fi
	done
	echo -e "${REWRITE}" 1>&2
	echo "$PASS"
}

_clevis() {
	if [ "$ACTION" == "luksunlock" ] ; then
		echo_ "Enable auto unlock LUKS partition using tpm2:"
		echo -e "  DEV: $LUKSDEV\n  UUID: $ROOTUUID\n  $(echo_ "Mount point:") /"
		getpass | clevis luks bind -d "$LUKSDEV" tpm2 '{"pcr_bank":"sha256","pcr_ids":"0,7"}' >> $LOG 2>&1
		if [ $? -ne 0 ] ; then
			echo_ "Bind device error!!!"
			EXITCODE=${LINENO}
			return
		fi
	elif [ "$ACTION" == "lukslock" ] ; then
		echo_ "Disable auto unlock"
		echo -e "  DEV: $LUKSDEV\n  UUID: $ROOTUUID\n $(echo_ "Mount point:") /"
		for SLOT in $(cryptsetup luksDump $LUKSDEV |sed -n '/clevis/{n;p;}' |awk '{print $NF}') ; do
			clevis luks unbind -d "$LUKSDEV" -s "$SLOT" -f >>$LOG 2>&1
			[ $? -eq 0 ] && UNBINDED=yes
		done
	[ "$UNBINDED" ] || echo_ "Unind device error"
	[ "$UNBINDED" ] || EXITCODE=${LINENO}
	fi
}

_systemd(){
	if [ "$ACTION" == "luksunlock" ] ; then
		echo_ "Enable auto unlock LUKS partition using tpm2:"
		echo -e "  DEV: $LUKSDEV\n  UUID: $ROOTUUID\n  $(echo_ "Mount point:") /"
		expect_crypenroll $LUKSDEV
		echo ''
		if [ $? -ne 0 ] ; then
			echo_ "Bind device error!!!"
			EXITCODE=${LINENO}
			return
		fi
	elif [ "$ACTION" == "lukslock" ] ; then
		echo_ "Disable auto unlock"
		echo -e "  DEV: $LUKSDEV\n  UUID: $ROOTUUID\n  $(echo_ "Mount point:") /"
		for SLOT in $(systemd-cryptenroll "$LUKSDEV" |grep tpm2 |awk '{print $1}') ; do
			systemd-cryptenroll --wipe-slot=$SLOT $LUKSDEV >>$LOG 2>&1
			[ $? -eq 0 ] && UNBINDED=yes
		done
	[ "$UNBINDED" ] || echo_ "Unbind device error"
	[ "$UNBINDED" ] || EXITCODE=${LINENO}
	fi
}

# send password to systemd-cryptenroll using python pexpect
# need to hide systemd-cryptenroll english messages
expect_crypenroll(){
	enroll=$(LC_ALL=C /usr/bin/python3 <(cat << END
import pexpect
try:
  proc = pexpect.spawn("systemd-cryptenroll --tpm2-device=auto --tpm2-pcrs=0+7 $1")
  proc.timeout = 10
  ret1 = proc.expect([ "Please enter current passphrase for disk", "already enrolled" ])
  if ret1 == 0:
    proc.sendline("$(getpass)")
    ret2 = proc.expect([ "not correct", "token enrolled" ])
    if ret2 == 1:
      print("OK")
    if ret2 == 0:
      print("wrong_pass")
  elif ret1 == 1:
    print("enrolled")
  
except(pexpect.EOF, pexpect.TIMEOUT):
  pass
END
)
)
	if [ "$enroll" == "wrong_pass" ] ; then
		echo_ "Wrong password!!!"  >> $LOG
		return 1
    elif [ "$enroll" == "enrolled" ] ; then
		echo_ "Aready enrolled"  >> $LOG
		return 2
	elif [ "$enroll" != "OK" ] ; then
		echo_ "Unknown error"  >> $LOG
		return 3
	fi
	return 0
}




agreement(){
	if ! [ -e /tmp/luksunlock_agree ] ; then
	    echo
	    echo_ "Attention! The application is potentially dangerous."
	    echo_ "Can corrupt the work of programs which using TPM2, including the Windows OS"
	    echo_ "If you use programs such a BitLocker,"
	    echo_ "make a backup copy of the encrypted data before using it."
	    echo
	    echo_ "Please enter 'Yes' to continue."
	    read ans
	    agree="${ans,,}"
	    if [ "$agree" == 'yes' ] || [ "$agree" == $(echo_ 'yes') ] ; then
		> /tmp/luksunlock_agree
	    else
		exit ${LINENO}
	    fi
	fi
	echo -e "${REWRITE}\n" 1>&2
}

is_swap(){
	local luks_dev swap_devs
	luks_dev=$(realpath /dev/mapper/$1)
	swap_devs=$(swapon --show=NAME |sed '1d')
	if [ -n "$swap_devs" ] && echo "$swap_devs" |grep -qw "$luks_dev"  ; then 
		echo "$luks_dev"
		return 0
	fi
	return ${LINENO} 
}

# auto unlock additional mounted LUKS partitions (/home, /opt, /var...)
# using key-file in root (/) LUKS partition
another_crypt(){
	local mapper mpoint dev UUIDS
	UUIDS=$(blkid -t TYPE=crypto_LUKS -s UUID |grep -v $LUKSDEV |sed -e 's:^.*UUID="::' -e 's/"$//')
	if [ "$ACTION" == "luksunlock" ] ; then
		[ $EXITCODE -eq 0 ] || return
		for a in $UUIDS ; do
			mapper=$(dmsetup table |grep $a |cut -f1 -d ':')
			[ "$mapper" ] || not_in_use $a || continue
			mpoint=$(grep $mapper /proc/mounts |cut -f2 -d' ')
			dev=$(blkid -t UUID=$a |cut -f1 -d':')
			[ ! -b "$dev" ] && continue
			echo ''
			echo_ 	"Enable auto unlock LUKS partition:"
			swap_dev=$(is_swap "$mapper")
			if [ -z "$mpoint" ] && [ -n "$swap_dev" ] ; then
				echo -e "  DEV: $dev\n  UUID: $a\n  SWAP: $swap_dev"
				# remove rd.luks.uuid=LUKS from grub configuration
				sed -i 's/rd.luks.uuid='$mapper'//' /etc/default/grub							|| EXITCODE=${LINENO}
				swapoff "$swap_dev" >>$LOG 2>&1
				echo_ "Reconfigure grub2. Please wait..."
				update-grub2 >>$LOG 2>&1
				mkswap "$swap_dev" >>$LOG 2>&1
				# swapon "$swap_dev" >>$LOG 2>&1
				# setup encrypted swap with random password
				sed -i '/.*'$a'.*/d' /etc/crypttab												|| EXITCODE=${LINENO}
				echo "luks-$a UUID=$a /dev/urandom \
cipher=aes-cbc-essiv:sha256,hash=ripemd160,size=256,swap " >> /etc/crypttab						|| EXITCODE=${LINENO}
				# fix fstab
				sed -i '/\bswap/d' /etc/fstab													|| EXITCODE=${LINENO}
				echo "/dev/mapper/luks-$a none swap \
defaults,x-systemd.device-timeout=30 0 0" >> /etc/fstab											|| EXITCODE=${LINENO}
			elif [ -n "$mpoint" ] ; then
				echo -e "  DEV: $dev\n  UUID: $a\n  $(echo_ "Mount point:") $mpoint"
				mkdir -p /etc/luksunlock/keys
				dd if=/dev/urandom of=/etc/luksunlock/keys/${a}.key bs=1024 count=4 >> $LOG 2>&1
				chown root:root /etc/luksunlock/keys/${a}.key
				chmod 0400 /etc/luksunlock/keys/${a}.key
				getpass | cryptsetup luksAddKey $dev /etc/luksunlock/keys/${a}.key >> $LOG 2>&1	|| EXITCODE=${LINENO}
				sed -i '/.*'$a'.*/d' /etc/crypttab												|| EXITCODE=${LINENO}
				echo "luks-$a UUID=$a  /etc/luksunlock/keys/${a}.key luks" >> /etc/crypttab	 	|| EXITCODE=${LINENO}
			else
				echo -e "DEV: $dev\nUUID: $a"
				echo_ "Skip: $dev"
			fi
		done
	elif [ "$ACTION" == "lukslock" ] ; then
		for a in $UUIDS ; do
			[ -f /etc/luksunlock/keys/${a}.key ] || continue
			is_swap "$mapper" && continue
		mapper=$(dmsetup table |grep $a |cut -f1 -d ':')
		mpoint=$(grep $mapper /proc/mounts |cut -f2 -d' ')
		dev=$(blkid -t UUID=$a |cut -f1 -d':')
		[ -b $dev ] || continue
		echo ''
		echo_ "Disable auto unlock LUKS partition:"
		echo -e "  DEV: $dev\n  UUID: $a\n  $(echo_ "Mount point:") $mpoint"
		cryptsetup luksRemoveKey $dev /etc/luksunlock/keys/${a}.key 		|| EXITCODE=${LINENO}
		sed -i '/.*'$a'.*/d' /etc/crypttab									|| EXITCODE=${LINENO}
		echo "luks-$a UUID=$a none discard"  >> /etc/crypttab	 			|| EXITCODE=${LINENO}
		rm -f /etc/luksunlock/keys/${a}.key
		done
	fi

}

# setup auto unlock for one of mounted luks partitions using tpm2
one_part_unlock(){
	local mapper mpoint dev UUIDS
	UUIDS=$(blkid -t TYPE=crypto_LUKS -s UUID |sed -e 's:^.*UUID="::' -e 's/"$//')
	if [ ! "$UUIDS" ] ; then
		echo_ "Luks partitions not found. Nothing to do!"
		exit
	fi
	if [ "$ACTION" == "luksunlock" ] ; then
		unset unlocked
		for a in $UUIDS ; do
			mapper=$(dmsetup table |grep $a |cut -f1 -d ':')
			[ "$mapper" ] || not_in_use $a || continue 
			mpoint=$(grep $mapper /proc/mounts |cut -f2 -d' ')
			dev=$(blkid -t UUID=$a |cut -f1 -d':')
			[ ! -b "$dev" ] || [ -z "$mpoint" ] && continue
			ask=$(echo_ "Do you want to automatically mount this LUKS partition using tpm2? (Y/N)")
			printf '%'$(( $(echo $ask |wc -m) -1 ))'s\n' |sed 's/ /=/g'
			echo -e "  DEV: $dev\n  UUID: $a\n  $(echo_ "Mount point:") $mpoint"
			echo "$ask"
			read ans
			if [ "$ans" == Y ] || [ "$ans" == y ] ; then
				expect_crypenroll $dev
				sed -i '/.*'$a'.*/d' /etc/crypttab								|| EXITCODE=${LINENO}
				echo "luks-$a UUID=$a  - tpm2-device=auto" >> /etc/crypttab	 	|| EXITCODE=${LINENO}
				printf '%'$(( $(echo $ask |wc -m) -1 ))'s\n' |sed 's/ /=/g'
				unlocked=yes
				break
			fi
		done
		[ "$unlocked" ] || EXITCODE=${LINENO}
	elif [ "$ACTION" == "lukslock" ] ; then
		unset locked
		for a in $UUIDS ; do
			grep -q ${a}'.*tpm2-device=auto' /etc/crypttab || not_in_use $a || continue
			mapper=$(dmsetup table |grep $a |cut -f1 -d ':')
			[ "$mapper" ]  || not_in_use $a || continue
			mpoint=$(grep $mapper /proc/mounts |cut -f2 -d' ')
			[ "$mpoint" ] || not_in_use $a || continue
			dev=$(blkid -t UUID=$a |cut -f1 -d':')
			[ -b "$dev" ] || not_in_use $a || continue
			ask=$(echo_ "Disable auto mount LUKS partition:")
			printf '%'$(( $(echo $ask |wc -m) -1 ))'s\n' |sed 's/ /=/g'
			echo $ask
			echo -e "  DEV: $dev\n  UUID: $a\n  $(echo_ "Mount point:") $mpoint"
			sed -i '/.*'$a'.*/d' /etc/crypttab		 							|| EXITCODE=${LINENO}
			echo "luks-$a UUID=$a none discard"  >> /etc/crypttab 				|| EXITCODE=${LINENO}
			for SLOT in $(systemd-cryptenroll "$dev" |grep tpm2 |awk '{print $1}') ; do
				systemd-cryptenroll --wipe-slot="$SLOT" "$dev" >>$LOG 2>&1 && locked='yes'
			done
			printf '%'$(( $(echo $ask |wc -m) -1 ))'s\n' |sed 's/ /=/g'
		done
		[ "$locked" ] || EXITCODE=${LINENO}
	fi
}

not_in_use(){
	echo -n "luks UUID: $1 - " >> $LOG
	echo_ "is not in use"  >> $LOG
	return 1
}

if [ "$1" ] && [ "$1" == '-h' -o "$1" == '--help' ] ; then
	echo -n "luksunlock, lukslock - "
	echo_ "Utils for automatic unlocking of LUKS system partitions using tpm2 module." 
	echo ''
	echo_ "Usage:"
	echo -n "luksunlock - "
	echo_ "Enable auto unlock"
	echo -n "lukslock - "
	echo_ "Disable auto unlock"
	echo ''
	echo_ "If the root partition (/) is encrypted, its decryption will be bound to tpm2,"
	echo_ "keys for the remaining LUKS partitions  will be created   in /etc/luksunlock/keys/."
	echo_ "If the root partition is not encrypted, only one LUKS partition of choice will be bound to tpm2." 
	exit
fi
# checks
if [ $(id -u) -ne 0 ] ; then
	echo_ "Must be root"
	exit ${LINENO}
fi

# exit if tpm2 not found
if ! ls /dev/tpm{0,rm0} >/dev/null 2>&1; then
	echo_ "Machine have no tpm2, or kernel tpm2 modules not active"
	exit ${LINENO}
fi

# choose from clevis and systemd
if rpm -q luksunlock-clevis >/dev/null 2>&1; then
	MODE='clevis'
elif rpm -q luksunlock-systemd >/dev/null 2>&1; then
	MODE='systemd'
else
	echo_ "luksunlock-clevis or luksunlock-systemd are not installed!"
	exit ${LINENO}
fi

# check that tpm2-abrmd daemon is running
if [ "$MODE" == 'systemd' ] && ! systemctl is-active tpm2-abrmd.service >/dev/null ; then
	systemctl enable tpm2-abrmd.service >/dev/null 2>&1
	echo_ "Enable tpm2-abrmd service"
	systemctl start tpm2-abrmd.service >/dev/null 2>&1 || 	exit ${LINENO}
fi

# find LUKS partition with system rootfs
if [ "$1" ] ; then
	LUKSDEV=$1
else
	LUKSDEV=$(lsblk -lnp  -o NAME,MAJ:MIN | grep $(dmsetup table $(df --output=source / \
			|tail -n1) 2>/dev/null |awk '{print $(NF-1)}') 2>/dev/null |cut -f1 -d ' ')
fi

# check luks partition
if cryptsetup isLuks "$LUKSDEV" 2>/dev/null ; then
	agreement
	ROOTUUID=$(blkid $LUKSDEV -s UUID -o value)
	_$MODE
	if [ "$EXITCODE" -eq 0 ] ; then
		echo_ "Generating init Ram Disk. Please wait..."
		initramfs-regen >>$LOG 2>&1 || EXITCODE="$?"
	fi
	another_crypt
else
	one_part_unlock
fi

if [ "$EXITCODE" -ne 0 ] ; then
	len=$(( $(cat "$LOG" |sed 's/./#/g' |sort |tail -n1 |wc -m) -1 ))
	printf '%'$len's\n' |sed 's/ /=/g'
	echo_ "Errors log:" 
	cat "$LOG"
	printf '%'$len's\n' |sed 's/ /=/g'
fi
rm -f "$LOG"
if [ "$EXITCODE" -eq 0 ] ; then
	[ "$ACTION" == "luksunlock" ] && echo && echo_ "Don't forget to set UEFI password !!!"
	echo_ "Complete!" && exit 0
fi

exit "$EXITCODE"
