#!/bin/bash
# Shell functions for most initialization scripts
. /etc/ash_functions

# Print <hidden> or <empty> depending on whether $1 is empty.  Useful to mask an
# optional password parameter.
mask_param() {
	if [ -z "$1" ]; then
		echo "<empty>"
	else
		echo "<hidden>"
	fi
}

# Trace a command with DEBUG, then execute it.
# A password parameter can be masked by passing --mask-position N before the
# command to execute, the debug trace will just indicate whether the password
# was empty or nonempty (which is important when use of a password is optional).
# N=0 is the name of the command to be executed, N=1 is its first parameter,
# etc.
DO_WITH_DEBUG() {
	if [ "$1" == "--mask-position" ]; then
		mask_position="$2"
		shift
		shift
		DEBUG_ARGS=("$@")

		DEBUG_ARGS[$mask_position]="$(mask_param "${DEBUG_ARGS[$mask_position]}")"
		DEBUG "${DEBUG_ARGS[@]}"
	else
		DEBUG "$@"
	fi
	"$@"
}

pcrs() {
	if [ "$CONFIG_TPM2_TOOLS" = "y" ]; then
		tpm2 pcrread sha256
	elif [ "$CONFIG_TPM" = "y" ]; then
		head -8 /sys/class/tpm/tpm0/pcrs
	fi
}

confirm_totp()
{
	TRACE "Under /etc/functions:confirm_totp"
	prompt="$1"
	last_half=X
	unset totp_confirm

	while true; do

		# update the TOTP code every thirty seconds
		date=`date "+%Y-%m-%d %H:%M:%S"`
		seconds=`date "+%s"`
		half=`expr \( $seconds % 60 \) / 30`
		if [ "$CONFIG_TPM" != "y" ]; then
			TOTP="NO TPM"
		elif [ "$half" != "$last_half" ]; then
			last_half=$half;
			TOTP=`unseal-totp` \
			|| recovery "TOTP code generation failed"
		fi

		echo -n "$date $TOTP: "

		# read the first character, non-blocking
		read \
			-t 1 \
			-n 1 \
			-s \
			-p "$prompt" \
			totp_confirm \
		&& break

		# nothing typed, redraw the line
		echo -ne '\r'
	done

	# clean up with a newline
	echo
}

# Enable USB storage (if not already enabled), and wait for storage devices to
# be detected.  If USB storage was already enabled, no wait occurs, this would
# have happened already when USB storage was enabled.
enable_usb_storage()
{
	if ! lsmod | grep -q usb_storage; then
		timeout=0
		echo "Scanning for USB storage devices..."
		insmod /lib/modules/usb-storage.ko >/dev/null 2>&1 \
		|| die "usb_storage: module load failed"
		while [[ $(list_usb_storage | wc -l) -eq 0 ]]; do 
			[[ $timeout -ge 8 ]] && break
			sleep 1
			timeout=$(($timeout+1))
		done
	fi
}

list_usb_storage()
{
	TRACE "Under /etc/functions:list_usb_storage"
	stat -c %N /sys/block/sd* 2>/dev/null | grep usb |
		cut -f1 -d ' ' |
		sed "s/[']//g" |
		while read b; do
			# Ignore devices of size 0, such as empty SD card
			# readers on laptops attached via USB.
			if [ "$(cat "$b/size")" -gt 0 ]; then
				echo "$b"
			fi
		done |
		sed "s|/sys/block|/dev|" |
		while read b; do
			# If the device has a partition table, ignore it and
			# include the partitions instead - even if the kernel
			# hasn't detected the partitions yet.  Such a device is
			# never usable directly, and this allows the "wait for
			# disks" loop in mount-usb to correctly wait for the
			# partitions.
			# This check: [ $(fdisk -l "$b" | wc -l) -eq 5 ]
			# covers the case of a device without partition table but
			# formatted as fat32, which contains a sortof partition table.
			# this causes fdisk to not print the invalid partition table
			# message and instead it'll print an empty table with header.
			# In both cases the output is 5 lines: 3 about device info,
			# 1 empty line and the 5th will be the table header or the
			# unvalid message.
			DISK_DATA=$(fdisk -l "$b")
			if echo "$DISK_DATA" | grep -q "doesn't contain a valid partition table" || [ $(echo "$DISK_DATA" | wc -l) -eq 5 ]; then
				# No partition table, include this device
				echo "$b"
			else
				# Has a partition table, include partitions
				ls -1 "$b"* | awk 'NR!=1 {print $0}'
			fi
		done
}

confirm_gpg_card()
{
	TRACE "Under /etc/functions:confirm_gpg_card"
	read \
		-n 1 \
		-p "Please confirm that your GPG card is inserted [Y/n]: " \
		card_confirm
	echo

	if [ "$card_confirm" != "y" \
		-a "$card_confirm" != "Y" \
		-a -n "$card_confirm" ] \
	; then
		die "gpg card not confirmed"
	fi

	# setup the USB so we can reach the GPG card
	enable_usb

	echo -e "\nVerifying presence of GPG card...\n"
	# ensure we don't exit without retrying
	errexit=$(set -o | grep errexit | awk '{print $2}')
	set +e
	gpg --card-status > /dev/null
	if [ $? -ne 0 ]; then
	  # prompt for reinsertion and try a second time
	  read -n1 -r -p \
	      "Can't access GPG key; remove and reinsert, then press Enter to retry. " \
	      ignored
	  # restore prev errexit state
	  if [ "$errexit" = "on" ]; then
	    set -e
	  fi
	  # retry card status
	  gpg --card-status > /dev/null \
	  	|| die "gpg card read failed"
	fi
	# restore prev errexit state
	if [ "$errexit" = "on" ]; then
	  set -e
	fi
}

# Prompt for an owner password if it is not already set in tpm_password.  Sets
# tpm_password.  Tools should optionally accept a TPM password on the command
# line, since some flows need it multiple times and only one prompt is ideal.
prompt_tpm_password() {
	if [ -n "$tpm_password" ]; then
		return 0;
	fi

	read -s -p "TPM Owner password: " tpm_password
	echo # new line after password prompt
}

# Prompt for a new owner password when resetting the TPM.  Returned in
# key_password.  The password must be 1-32 characters and must be entered twice,
# the script will loop until this is met.
prompt_new_owner_password() {
	local key_password2
	key_password=1
	key_password2=2
	while [ "$key_password" != "$key_password2" ] || [ "${#key_password}" -gt 32 ] || [ -z "$key_password" ]; do
		read -s -p "New TPM owner passphrase (2 words suggested, 1-32 characters max): " key_password
		echo

		read -s -p "Repeat chosen TPM owner passphrase: " key_password2
		echo

		if [ "$key_password" != "$key_password2" ]; then
			echo "Passphrases entered do not match. Try again!"
			echo
		fi
	done
}

check_tpm_counter()
{
  TRACE "Under /etc/functions:check_tpm_counter"
  LABEL=${2:-3135106223}
  tpm_password="$3"
	# if the /boot.hashes file already exists, read the TPM counter ID
	# from it.
	if [ -r "$1" ]; then
		TPM_COUNTER=`grep counter- "$1" | cut -d- -f2`
	else
		warn "$1 does not exist; creating new TPM counter"
		prompt_tpm_password
		tpmr counter_create \
			-pwdo "$tpm_password" \
			-pwdc '' \
			-la $LABEL \
		| tee /tmp/counter \
		|| die "Unable to create TPM counter"
		TPM_COUNTER=`cut -d: -f1 < /tmp/counter`
	fi

	if [ -z "$TPM_COUNTER" ]; then
		die "$1: TPM Counter not found?"
	fi
}

read_tpm_counter()
{
	TRACE "Under /etc/functions:read_tpm_counter"
	tpmr counter_read -ix "$1" | tee "/tmp/counter-$1" \
	|| die "Counter read failed"
}

increment_tpm_counter()
{
	TRACE "Under /etc/functions:increment_tpm_counter"
	tpmr counter_increment -ix "$1" -pwdc '' \
		| tee /tmp/counter-$1 \
	|| die "Counter increment failed"
}

check_config() {
	TRACE "Under /etc/functions:check_config"
	if [ ! -d /tmp/kexec ]; then
		mkdir /tmp/kexec \
		|| die 'Failed to make kexec tmp dir'
	else
		rm -rf /tmp/kexec/* \
		|| die 'Failed to empty kexec tmp dir'
	fi

	if [ ! -r $1/kexec.sig -a "$CONFIG_BASIC" != "y" ]; then
		return
	fi

	if [ `find $1/kexec*.txt | wc -l` -eq 0 ]; then
		return
	fi

	if [ "$2" != "force" ]; then
		if ! sha256sum `find $1/kexec*.txt` | gpgv $1/kexec.sig - ; then
			die 'Invalid signature on kexec boot params'
		fi
	fi

	echo "+++ Found verified kexec boot params"
	cp $1/kexec*.txt /tmp/kexec \
	|| die "Failed to copy kexec boot params to tmp"
}

# Replace a file in a ROM (add it if the file does not exist)
replace_rom_file() {
	ROM="$1"
	ROM_FILE="$2"
	NEW_FILE="$3"
	
	if (cbfs.sh -o "$ROM" -l | grep -q "$ROM_FILE") then
		cbfs.sh -o "$ROM" -d "$ROM_FILE"
	fi
	cbfs.sh -o "$ROM" -a "$ROM_FILE" -f "$NEW_FILE"
}

replace_config() {
	TRACE "Under /etc/functions:replace_config"
	CONFIG_FILE=$1
	CONFIG_OPTION=$2
	NEW_SETTING=$3

	touch $CONFIG_FILE
# first pull out the existing option from the global config and place in a tmp file
	awk "gsub(\"^export ${CONFIG_OPTION}=.*\",\"export ${CONFIG_OPTION}=\\\"${NEW_SETTING}\\\"\")" /tmp/config > ${CONFIG_FILE}.tmp
	awk "gsub(\"^${CONFIG_OPTION}=.*\",\"${CONFIG_OPTION}=\\\"${NEW_SETTING}\\\"\")" /tmp/config >> ${CONFIG_FILE}.tmp

# then copy any remaining settings from the existing config file, minus the option you changed
	grep -v "^export ${CONFIG_OPTION}=" ${CONFIG_FILE} | grep -v "^${CONFIG_OPTION}=" >> ${CONFIG_FILE}.tmp || true
  sort ${CONFIG_FILE}.tmp | uniq > ${CONFIG_FILE}
	rm -f ${CONFIG_FILE}.tmp
}

# Set a config variable in a specific file to a given value - replace it if it
# exists, or add it.  If added, the variable will be exported.
set_config() {
	CONFIG_FILE="$1"
	CONFIG_OPTION="$2"
	NEW_SETTING="$3"
	
	if grep -q "$CONFIG_OPTION" "$CONFIG_FILE"; then
		replace_config "$CONFIG_FILE" "$CONFIG_OPTION" "$NEW_SETTING"
	else
		echo "export $CONFIG_OPTION=\"$NEW_SETTING\"" >>"$CONFIG_FILE"
	fi
}

# Set a value in config.user, re-combine configs, and update configs in the
# environment.
set_user_config() {
	CONFIG_OPTION="$1"
	NEW_SETTING="$2"

	set_config /etc/config.user "$CONFIG_OPTION" "$NEW_SETTING"
	combine_configs
	. /tmp/config
}

# Load a config value to a variable, defaulting to empty.  Does not fail if the
# config is not set (since it would expand to empty by default).
load_config_value()
{
  local config_name="$1"
  if grep -q "$config_name=" /tmp/config; then
    grep "$config_name=" /tmp/config | tail -n1 | cut -f2 -d '=' | tr -d '"'
  fi
}

# Generate a secret for TPM-less HOTP by reading the ROM.  Output is the
# sha256sum of the ROM (binary, not printable), which can be truncated to the
# supported secret length.
secret_from_rom_hash() {
	local ROM_IMAGE="/tmp/coreboot-notpm.rom"

	echo -e "\nTPM not detected; measuring ROM directly\n" 1>&2

	# Read the ROM if we haven't read it yet
	if [ ! -f "${ROM_IMAGE}" ]; then
		flash.sh -r "${ROM_IMAGE}" >/dev/null 2>&1 || return 1
	fi

	sha256sum "${ROM_IMAGE}" | cut -f1 -d ' ' | fromhex_plain
}

update_checksums()
{
	TRACE "Under /etc/functions:update_checksums"
	# ensure /boot mounted
	if ! grep -q /boot /proc/mounts ; then
		mount -o ro /boot \
		|| recovery "Unable to mount /boot"
	fi

	# remount RW
	mount -o rw,remount /boot

	# sign and auto-roll config counter
	extparam=
	if [ "$CONFIG_TPM" = "y" ];then
		if [ "$CONFIG_IGNORE_ROLLBACK" != "y" ]; then
			extparam=-r
		fi
	fi
	if ! kexec-sign-config -p /boot -u $extparam ; then
		rv=1
	else
		rv=0
	fi

	# switch back to ro mode
	mount -o ro,remount /boot

	return $rv
}

print_tree() {
	TRACE "Under /etc/functions:print_tree"
	find ./ ! -path './kexec*' -print0 | sort -z
}

# Escape zero-delimited standard input to safely display it to the user in e.g.
# `whiptail`, `less`, `echo`, `cat`. Doesn't produce shell-escaped output.
# Most printable characters are passed verbatim (exception: \).
# These escapes are used to replace their corresponding characters: #n#r#t#v#b
# Other characters are rendered as hexadecimal escapes.
# escape_zero [prefix] [escape character]
# prefix: \0 in the input will result in \n[prefix]
# escape character: character to use for escapes (default: #); \ may be interpreted by `whiptail`
escape_zero() {
	local prefix="$1"
	local echar="${2:-#}"
	local todo=""
	local echar_hex="$(echo -n "$echar" | xxd -p -c1)"
	[ ${#echar_hex} -eq 2 ] || die "Invalid escape character $echar passed to escape_zero(). Programming error?!"

	echo -e -n "$prefix"
	xxd -p -c1 | tr -d '\n' |
	{
		while IFS= read -r -n2 -d '' ; do
			if [ -n "$todo" ] ; then
				#REPLY == "  " is EOF
				[[ "$REPLY" == "  " ]] && echo '' || echo -e -n "$todo"
				todo=""
			fi

			case "$REPLY" in
				00)
					todo="\n$prefix"
					;;
				08)
					echo -n "${echar}b"
					;;
				09)
					echo -n "${echar}t"
					;;
				0a)
					echo -n "${echar}n"
					;;
				0b)
					echo -n "${echar}v"
					;;
				0d)
					echo -n "${echar}r"
					;;
				"$echar_hex")
					echo -n "$echar$echar"
					;;
				#interpreted characters:
				2[0-9a-f]|3[0-9a-f]|4[0-9a-f]|5[0-9abd-f]|6[0-9a-f]|7[0-9a-e])
					echo -e -n '\x'"$REPLY"
					;;
				# All others are escaped
				*)
					echo -n "${echar}x$REPLY"
					;;
			esac
		done
	}
}

# Currently heads doesn't support signing file names with certain characters
# due to https://bugs.busybox.net/show_bug.cgi?id=14226. Also, certain characters
# may be intepreted by `whiptail`, `less` et al (e.g. \n, \b, ...).
assert_signable() {
	TRACE "Under /etc/functions:assert_signable"
	# ensure /boot mounted
	if ! grep -q /boot /proc/mounts ; then
		mount -o ro /boot || die "Unable to mount /boot"
	fi

	find /boot -print0 > /tmp/signable.ref
	local del='\001-\037\134\177-\377'
	LC_ALL=C tr -d "$del" < /tmp/signable.ref > /tmp/signable.del || die "Failed to execute tr."
	if ! cmp -s "/tmp/signable.ref" "/tmp/signable.del" &> /dev/null ; then
		local user_out="/tmp/hash_output_mismatches"
		local add="Please investigate!"
		[ -f "$user_out" ] && add="Please investigate the following relative paths to /boot (where # are sanitized invalid characters):"$'\n'"$(cat "$user_out")"
		recovery "Some /boot file names contain characters that are currently not supported by heads: $del"$'\n'"$add"
	fi
	rm -f /tmp/signable.*
}

verify_checksums()
{
	TRACE "Under /etc/functions:verify_checksums"
	local boot_dir="$1"
	local gui="${2:-y}"

	(
		set +e -o pipefail
		local ret=0
		cd "$boot_dir" || ret=1
		sha256sum -c "$TMP_HASH_FILE" > /tmp/hash_output || ret=1

		# also make sure that the file & directory structure didn't change
		# (sha256sum won't detect added files)
		print_tree > /tmp/tree_output || ret=1
		if ! cmp -s "$TMP_TREE_FILE" /tmp/tree_output &> /dev/null ; then
			ret=1
			[[ "$gui" != "y" ]] && exit "$ret"
			# produce a diff that can safely be presented to the user
			# this is relatively hard as file names may e.g. contain backslashes etc.,
			# which are interpreted by whiptail, less, ...
			escape_zero "(new) " < "$TMP_TREE_FILE" > "${TMP_TREE_FILE}.user"
			escape_zero "(new) " < /tmp/tree_output > /tmp/tree_output.user
			diff "${TMP_TREE_FILE}.user" /tmp/tree_output.user | grep -E '^\+\(new\).*$' | sed -r 's/^\+\(new\)/(new)/g' >> /tmp/hash_output
			rm -f "${TMP_TREE_FILE}.user"
			rm -f /tmp/tree_output.user
		fi
		exit $ret
	)
	return $?
}

# detect and set /boot device
# mount /boot if successful
detect_boot_device()
{
	TRACE "Under /etc/functions:detect_boot_device"
	# unmount /boot to be safe
	cd / && umount /boot 2>/dev/null

	# check $CONFIG_BOOT_DEV if set/valid
	if [ -e "$CONFIG_BOOT_DEV" ]; then
		if mount -o ro $CONFIG_BOOT_DEV /boot >/dev/null 2>&1; then
			if ls -d /boot/grub* >/dev/null 2>&1; then
				# CONFIG_BOOT_DEV is valid device and contains an installed OS
				return 0
			fi
		fi
	fi

	# generate list of possible boot devices
	fdisk -l | grep "Disk /dev/" | cut -f2 -d " " | cut -f1 -d ":" > /tmp/disklist

	# filter out extraneous options
	> /tmp/boot_device_list
	for i in `cat /tmp/disklist`; do
		# remove block device from list if numeric partitions exist, since not bootable
		DEV_NUM_PARTITIONS=$((`ls -1 $i* | wc -l`-1))
		if [ ${DEV_NUM_PARTITIONS} -eq 0 ]; then
			echo $i >> /tmp/boot_device_list
		else
			ls $i* | tail -${DEV_NUM_PARTITIONS} >> /tmp/boot_device_list
		fi
	done

	# iterate thru possible options and check for grub dir
	for i in `cat /tmp/boot_device_list`; do
		umount /boot 2>/dev/null
		if mount -o ro $i /boot >/dev/null 2>&1; then
			if ls -d /boot/grub* >/dev/null 2>&1; then
				CONFIG_BOOT_DEV="$i"
				return 0
			fi
		fi
	done

	# no valid boot device found
	echo "Unable to locate /boot files on any mounted disk"
	umount /boot 2>/dev/null
	return 1
}

scan_boot_options()
{
	local bootdir config option_file
	bootdir="$1"
	config="$2"
	option_file="$3"

	if [ -r $option_file ]; then rm $option_file; fi
	for i in `find $bootdir -name "$config"`; do
		DO_WITH_DEBUG kexec-parse-boot "$bootdir" "$i" >> $option_file
	done
	# FC29/30+ may use BLS format grub config files
	# https://fedoraproject.org/wiki/Changes/BootLoaderSpecByDefault
	# only parse these if $option_file is still empty
	if [ ! -s $option_file ] && [ -d "$bootdir/loader/entries" ]; then
		for i in `find $bootdir -name "$config"`; do
			kexec-parse-bls "$bootdir" "$i" "$bootdir/loader/entries" >> $option_file
		done
	fi
}

calc()
{ 
	awk "BEGIN { print "$*" }"; 
}

# truncate a file to a size only if it is longer (busybox truncate lacks '<' and
# always sets the file size)
truncate_max_bytes() {
	local bytes="$1"
	local file="$2"
	if [ "$(stat -c %s "$file")" -gt "$bytes" ]; then
		truncate -s "$bytes" "$file"
	fi
}

# Busybox xxd -p pads the last line with spaces to 60 columns, which not only
# trips up many scripts, it's very difficult to diagnose by looking at the
# output.  Delete line breaks and spaces to really get plain hex output.
tohex_plain() {
	xxd -p | tr -d '\n '
}

# Busybox xxd -p -r silently truncates lines longer than 60 hex chars.
# Shorter lines are OK, spaces are OK, and even splitting a byte across lines is
# allowed, so just fold the text to maximum 60 column lines.
# Note that also unlike GNU xxd, non-hex chars in input corrupt the output (GNU
# xxd ignores them).
fromhex_plain() {
	fold -w 60 | xxd -p -r
}

print_battery_health()
{
	if [ -d /sys/class/power_supply/BAT* ]; then
        	battery_health=$(calc $(cat /sys/class/power_supply/BAT*/charge_full)/$(cat /sys/class/power_supply/BAT*/charge_full_design)*100 | awk -F "." {'print $1'})
		echo "$battery_health"
	fi
}

print_battery_charge()
{
	if [ -d /sys/class/power_supply/BAT* ]; then
                battery_charge=$(calc $(cat /sys/class/power_supply/BAT*/charge_now)/$(cat /sys/class/power_supply/BAT*/charge_full)*100 | awk -F "." {'print $1'})
                echo "$battery_charge"
        fi
}

generate_random_mac_address()
{
	#Borrowed from https://stackoverflow.com/questions/42660218/bash-generate-random-mac-address-unicast
	hexdump -n 6 -ve '1/1 "%.2x "' /dev/urandom | awk -v a="2,6,a,e" -v r="$RANDOM" 'BEGIN{srand(r);}NR==1{split(a,b,",");r=int(rand()*4+1);printf "%s%s:%s:%s:%s:%s:%s\n",substr($1,0,1),b[r],$2,$3,$4,$5,$6}'
}

# Add a command to be invoked at exit.  (Note that trap EXIT replaces any
# existing handler.)  Commands are invoked in reverse order, so they can be used
# to clean up resources, etc.
# The parameters are all executed as-is and do _not_ require additional quoting
# (unlike trap).  E.g.:
# at_exit shred "$file" #<-- file is expanded when calling at_exit, no extra quoting needed
at_exit() {
	AT_EXIT_HANDLERS+=("$@") # Command and args
	AT_EXIT_HANDLERS+=("$#") # Number of elements in this command
}

# Array of all exit handler command arguments with lengths of each command at
# the end.  For example:
#   at_exit echo hello
#   at_exit echo a b c
# results in:
# AT_EXIT_HANDLERS=(echo hello 2 echo a b c 4)

AT_EXIT_HANDLERS=()
# Each handler is an array AT_EXIT_HANDLER_{i}
run_at_exit_handlers() {
	local cmd_pos cmd_len
	cmd_pos="${#AT_EXIT_HANDLERS[@]}"
	# Silence trace if there are no handlers, this is common and occurs a lot
	[ "$cmd_pos" -gt 0 ] && DEBUG "Running at_exit handlers"
	while [ "$cmd_pos" -gt 0 ]; do
		cmd_pos="$((cmd_pos-1))"
		cmd_len="${AT_EXIT_HANDLERS[$cmd_pos]}"
		cmd_pos="$((cmd_pos-cmd_len))"
		"${AT_EXIT_HANDLERS[@]:$cmd_pos:$cmd_len}"
	done
}
trap run_at_exit_handlers EXIT