#!/bin/sh
# Automated setup of TPM, GPG keys, and disk

set -o pipefail

# use TERM to exit on error
trap "exit 1" TERM
export TOP_PID=$$

## Static local variables

CLEAR="--clear"
CONTINUE="--yes-button Continue"
CANCEL="--no-button Cancel"
HEIGHT="150"
WIDTH="220"

USER_PIN_DEF=123456
ADMIN_PIN_DEF=12345678
TPM_PASS_DEF=12345678
CUSTOM_PASS=""

## External files sourced

. /etc/functions
. /tmp/config

## functions

die() {

	local msg=$1
	if [ -n "$msg" ]; then
	    echo -e "\n$msg"
	fi
	kill -s TERM $TOP_PID
	exit 1
}

whiptail_error() 
{
    local msg=$1
    if [ "$msg" = "" ]; then
        die "whiptail error: An error msg is required"
    fi
    whiptail --msgbox "${msg}\n\n" $WIDTH $HEIGHT $CONFIG_ERROR_BG_COLOR --title "Error"
}

whiptail_error_die() 
{
    whiptail_error "$@"
    die
}

gpg_key_reset()
{
    GPG_KEY_NAME=`date +%Y%m%d%H%M%S`
    # Factory reset GPG card
    {
        echo admin
        echo factory-reset
        echo y
        echo yes
    } | gpg --command-fd=0 --status-fd=1 --pinentry-mode=loopback --card-edit \
        > /tmp/gpg_card_edit_output 2>/dev/null
    if [ $? -ne 0 ]; then
        ERROR=`cat /tmp/gpg_card_edit_output`
        whiptail_error_die "GPG Key factory reset failed!\n\n$ERROR"
    fi
    # Generate OEM GPG keys
    {
        echo admin
        echo generate
        echo n
        echo ${ADMIN_PIN_DEF}
        echo ${USER_PIN_DEF}
        echo 0
        echo y
        echo "OEM Key"
        echo "oem-${GPG_KEY_NAME}@example.com"
        echo "OEM-generated key"
    } | gpg --command-fd=0 --status-fd=2 --pinentry-mode=loopback --card-edit \
        > /tmp/gpg_card_edit_output 2>/dev/null 
    if [ $? -ne 0 ]; then
        ERROR=`cat /tmp/gpg_card_edit_output`
        whiptail_error_die "GPG Key automatic keygen failed!\n\n$ERROR"
    fi
}
gpg_key_change_pin()
{
    # 1 = user PIN, 3 = admin PIN
    PIN_TYPE=$1
    PIN_ORIG=$2
    PIN_NEW=$3
    # Change PIN
    {
        echo admin
        echo passwd
        echo ${PIN_TYPE}
        echo ${PIN_ORIG}
        echo ${PIN_NEW}
        echo ${PIN_NEW}
        echo q
        echo q
    } | gpg --command-fd=0 --status-fd=2 --pinentry-mode=loopback --card-edit \
        > /tmp/gpg_card_edit_output 2>/dev/null 
    if [ $? -ne 0 ]; then
        ERROR=`cat /tmp/gpg_card_edit_output`
        whiptail_error_die "GPG Key PIN change failed!\n\n$ERROR"
    fi
}

generate_checksums()
{
    # ensure /boot mounted
    if ! grep -q /boot /proc/mounts ; then
        mount -o rw /boot || whiptail_error_die "Unable to mount /boot"
    else
        mount -o remount,rw /boot || whiptail_error_die "Unable to mount /boot"
    fi

    # clear any existing checksums/signatures
    rm /boot/kexec* 2>/dev/null

    # create Heads TPM counter
    tpm counter_create \
        -pwdo "$TPM_PASS_DEF" \
        -pwdc '' \
        -la -3135106223 \
        | tee /tmp/counter \
        || whiptail_error_die "Unable to create TPM counter"
    TPM_COUNTER=`cut -d: -f1 < /tmp/counter`

    # increment TPM counter
    increment_tpm_counter $TPM_COUNTER >/dev/null 2>&1 \
        || whiptail_error_die "Unable to increment tpm counter"

    # create rollback file
    sha256sum /tmp/counter-$TPM_COUNTER > /boot/kexec_rollback.txt 2>/dev/null \
        || whiptail_error_die "Unable to create rollback file"

    # set default boot option
    set_default_boot_option

    # generate hashes
    find /boot -type f ! -name '*kexec*' \
        | xargs sha256sum > /boot/kexec_hashes.txt 2>/dev/null \
        || whiptail_error_die "Error generating kexec hashes"

    param_files=`find /boot/kexec*.txt`
    [ -z "$param_files" ] \
        && whiptail_error_die "No kexec parameter files to sign"

    # sign kexec boot files
    if sha256sum $param_files 2>/dev/null | gpg \
            --pinentry-mode loopback \
            --passphrase "$USER_PIN_DEF" \
            --digest-algo SHA256 \
            --detach-sign \
            -a \
            > /boot/kexec.sig 2>/tmp/error; then
        # successful - update the validated params
        if ! check_config /boot >/dev/null 2>/tmp/error ; then
            cat /tmp/error
            ret=1
        else
            ret=0
        fi
    else
        cat /tmp/error
        ret=1
    fi

    # done writing to /boot, switch back to RO
    mount -o ro,remount /boot

    if [ $ret = 1 ] ; then
        ERROR=$(tail -n 1 /tmp/error)
        whiptail_error_die "Error signing kexec boot files:\n\n$ERROR"
    fi
}

set_default_boot_option()
{
    option_file="/tmp/kexec_options.txt"
    tmp_menu_file="/tmp/kexec/kexec_menu.txt"
    hash_file="/boot/kexec_default_hashes.txt"

    mkdir -p /tmp/kexec/
    rm $option_file 2>/dev/null
    # parse boot options from grub.cfg
    for i in `find /boot -name "grub.cfg"`; do
        kexec-parse-boot "/boot" "$i" >> $option_file
    done
    [ ! -r $option_file ] \
        && whiptail_error_die "Failed to parse any boot options"

    # sort boot options
    sort -r $option_file | uniq > $tmp_menu_file

    ## save first option as default
    entry=`head -n 1 $tmp_menu_file | tail -1`

    # clear existing default configs
    rm "/boot/kexec_default.*.txt" 2>/dev/null

    # write new config
    echo "$entry" > /boot/kexec_default.1.txt

    # validate boot option
    cd /boot && /bin/kexec-boot -b "/boot" -e "$entry" -f \
        | xargs sha256sum > $hash_file 2>/dev/null \
        || whiptail_error_die "Failed to create hashes of boot files"
}

## main script start

# check for args
if [ "$1" != "" ]; then
    title_text=$1
else
    title_text="OEM Factory Reset"
fi
if [ "$2" != "" ]; then
    bg_color=$2
else
    bg_color=""
fi

# show warning prompt
if ! whiptail --yesno "
    This operation will automatically:\n\n
        * ERASE the TPM and reset it with a default password\n
        * ERASE any keys or passwords on the GPG smart card,\n
        reset it to a factory state, and generate new keys\n
        * Add the new GPG key to the firmware and reflash it\n
        * Sign all of the files in /boot with the new GPG key\n\n
    It requires that you already have an OS installed on a\n
    dedicated /boot partition. Do you wish to continue?\n" \
        $WIDTH $HEIGHT $CONTINUE $CANCEL $CLEAR $bg_color --title "$title_text" ; then
    exit 1
fi

# Prompt to change default passwords
echo -e -n "Would you like to set a custom password? [y/N]: "
read -n 1 prompt_output
echo
if [ "$prompt_output" == "y" \
  -o "$prompt_output" == "Y" ] \
; then
  echo -e "\nThe custom password will be used for the
TPM admin and GPG user/admin passwords.
It must be at least 8 characters in length.\n"
  CUSTOM_PASS=""
  echo
  while [[  ${#CUSTOM_PASS} -lt 8 ]] ; do
    echo -e -n "Enter the custom password: "
    read CUSTOM_PASS
  done
  echo
  TPM_PASS_DEF=$CUSTOM_PASS
fi

## sanity check the USB, GPG key, and boot device before proceeding further

# mount USB, then remount rw
echo -e "\nChecking for USB media...\n"
# ensure /media not mounted
umount /media 2>/dev/null
# mount-usb will detect and prompt if no USB inserted
if ! mount-usb rw 2>/tmp/error; then
    ERROR=$(tail -n 1 /tmp/error)
    whiptail_error_die "Unable to mount USB on /media:\n\n${ERROR}"
fi

# ensure GPG key connected
echo -e "\nChecking for GPG Key...\n"
# USB kernel modules already loaded via mount-usb
if ! gpg --card-status >/dev/null 2>&1 ; then
    whiptail_error "Can't access GPG Key; remove and reinsert, then press Enter to retry."
    if ! gpg --card-status >/dev/null 2>/tmp/error ; then
        ERROR=$(tail -n 1 /tmp/error)
        whiptail_error_die "Unable to detect GPG Key:\n\n${ERROR}"
    fi
fi

# detect and set /boot device
echo -e "\nDetecting and setting boot device...\n"
if ! detect_boot_device ; then
  whiptail_error_die "Unable to locate /boot files on any mounted disk"
else
  echo -e "Boot device set to $CONFIG_BOOT_DEV\n"
fi

# update configs
replace_config /etc/config.user "CONFIG_BOOT_DEV" "$CONFIG_BOOT_DEV"
combine_configs

## reset TPM and set default password
echo -e "\nResetting TPM...\n"
{
    echo $TPM_PASS_DEF
    echo $TPM_PASS_DEF
} | /bin/tpm-reset >/dev/null 2>/tmp/error
if [ $? -ne 0 ]; then
    ERROR=$(tail -n 1 /tmp/error)
    whiptail_error_die "Error resetting TPM:\n\n${ERROR}"
fi

# clear local keyring
rm /.gnupg/*.gpg 2>/dev/null
rm /.gnupg/*.kbx 2>/dev/null
gpg --list-keys >/dev/null 2>&1

## reset the GPG Key
echo -e "\nResetting GPG Key...\n(this will take a minute or two)\n"
gpg_key_reset

if [ "$CUSTOM_PASS" != "" ]; then
  echo -e "\nChanging default GPG Admin PIN\n"
  gpg_key_change_pin "3" "$ADMIN_PIN_DEF" "$CUSTOM_PASS"
  echo -e "\nChanging default GPG User PIN\n"
  gpg_key_change_pin "1" "$USER_PIN_DEF" "$CUSTOM_PASS"
  USER_PIN_DEF=$CUSTOM_PASS
  ADMIN_PIN_DEF=$CUSTOM_PASS
fi

## export generated key to USB
echo -e "\nExporting generated key to USB...\n"
# parse name of generated key
GPG_GEN_KEY=`grep -A1 pub /tmp/gpg_card_edit_output | tail -n1 | sed -nr 's/^([ ])*//p'`
PUBKEY="/tmp/${GPG_GEN_KEY}.asc"
# export pubkey to file
if ! gpg --export --armor $GPG_GEN_KEY > "${PUBKEY}" 2>/tmp/error ; then
    ERROR=$(tail -n 1 /tmp/error)
    whiptail_error_die "GPG Key gpg export to file failed!\n\n$ERROR"
fi
# copy to USB
if ! cp "${PUBKEY}" "/media/${GPG_GEN_KEY}.asc" 2>/tmp/error ; then
    ERROR=$(tail -n 1 /tmp/error)
    whiptail_error_die "Key export error: unable to copy ${GPG_GEN_KEY}.asc to /media:\n\n$ERROR"
fi
umount /media 2>/dev/null

## flash generated key to ROM
echo -e "\nReading current firmware...\n(this will take a minute or two)\n"
/bin/flash.sh -r /tmp/oem-setup.rom >/dev/null 2>/tmp/error
if [ ! -s /tmp/oem-setup.rom ]; then
    ERROR=$(tail -n 1 /tmp/error)
    whiptail_error_die "Error reading current firmware:\n\n$ERROR"
fi

# ensure key imported locally
if ! cat "$PUBKEY" | gpg --import >/dev/null 2>/tmp/error ; then
    ERROR=$(tail -n 1 /tmp/error)
    whiptail_error_die "Error importing GPG key:\n\n$ERROR"
fi
# update /.gnupg/trustdb.gpg to ultimately trust all user provided public keys
if ! gpg --list-keys --fingerprint --with-colons 2>/dev/null \
        | sed -E -n -e 's/^fpr:::::::::([0-9A-F]+):$/\1:6:/p' \
        | gpg --import-ownertrust >/dev/null 2>/tmp/error ; then
    ERROR=$(tail -n 1 /tmp/error)
    whiptail_error_die "Error importing GPG ownertrust:\n\n$ERROR"
fi
if ! gpg --update-trust >/dev/null 2>/tmp/error ; then
    ERROR=$(tail -n 1 /tmp/error)
    whiptail_error_die "Error updating GPG ownertrust:\n\n$ERROR"
fi
# clear any existing heads/gpg files from current firmware
for i in `cbfs -o /tmp/oem-setup.rom -l | grep -e "heads/"`; do
    cbfs -o /tmp/oem-setup.rom -d $i
done
# add heads/gpg files to current firmware
if [ -e /.gnupg/pubring.kbx ];then
    cbfs -o /tmp/oem-setup.rom -a "heads/initrd/.gnupg/pubring.kbx" -f /.gnupg/pubring.kbx
    if [ -e /.gnupg/pubring.gpg ];then
        rm /.gnupg/pubring.gpg
    fi
elif [ -e /.gnupg/pubring.gpg ];then
    cbfs -o /tmp/oem-setup.rom -a "heads/initrd/.gnupg/pubring.gpg" -f /.gnupg/pubring.gpg
fi
if [ -e /.gnupg/trustdb.gpg ]; then
    cbfs -o /tmp/oem-setup.rom -a "heads/initrd/.gnupg/trustdb.gpg" -f /.gnupg/trustdb.gpg
fi
# persist user config changes (boot device)
if [ -e /etc/config.user ]; then
    cbfs -o /tmp/oem-setup.rom -a "heads/initrd/etc/config.user" -f /etc/config.user
fi
# flash updated firmware image
echo -e "\nAdding generated key to current firmware and re-flashing...\n"
if ! /bin/flash.sh /tmp/oem-setup.rom >/dev/null 2>/tmp/error ; then
    ERROR=$(tail -n 1 /tmp/error)
    whiptail_error_die "Error flashing updated firmware image:\n\n$ERROR"
fi

## sign files in /boot and generate checksums
echo -e "\nSigning boot files and generating checksums...\n"
generate_checksums

## all done -- reboot
whiptail --msgbox "
    The OEM Factory Reset has completed successfully\n\n
    After rebooting, you will need to generate new TOTP/HOTP secrets\n
    when prompted in order to complete the setup process.\n\n
    Press any key to reboot.\n" \
    $WIDTH $HEIGHT --title "OEM Factory Reset Complete"

reboot