#!/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="30"
WIDTH="90"

USER_PIN_DEF=123456
ADMIN_PIN_DEF=12345678
TPM_PASS_DEF=12345678
USER_PIN=""
ADMIN_PIN=""
TPM_PASS=""

# What are the Security components affected by custom passwords
CUSTOM_PASS_AFFECTED_COMPONENTS=""

if [ "$CONFIG_TPM" = "y" ]; then
    CUSTOM_PASS_AFFECTED_COMPONENTS="TPM Ownership password"
fi
CUSTOM_PASS_AFFECTED_COMPONENTS="
$CUSTOM_PASS_AFFECTED_COMPONENTS
GPG Admin PIN
GPG User PIN"

RSA_KEY_LENGTH=3072

GPG_USER_NAME="OEM Key"
GPG_KEY_NAME=`date +%Y%m%d%H%M%S`
GPG_USER_MAIL="oem-${GPG_KEY_NAME}@example.com"
GPG_USER_COMMENT="OEM-generated key"
SKIP_BOOT="n"

## 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 $BG_COLOR_ERROR --msgbox "${msg}\n\n" $HEIGHT $WIDTH $BG_COLOR_ERROR --title "Error"
}

whiptail_error_die() 
{
    whiptail_error "$@"
    die
}

gpg_key_reset()
{
    # 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
    # If Nitrokey Storage is inserted, reset AES keys as well
    if lsusb | grep -q "20a0:4109" &&  [ -x /bin/hotp_verification ] ; then
        /bin/hotp_verification regenerate ${ADMIN_PIN_DEF}
    fi
    # Set RSA key length
    {
        echo admin
        echo key-attr
        echo 1 # RSA
        echo ${RSA_KEY_LENGTH} #Signing key size set to RSA_KEY_LENGTH
        echo ${ADMIN_PIN_DEF}
        echo 1 # RSA
        echo ${RSA_KEY_LENGTH} #Encryption key size set to RSA_KEY_LENGTH
        echo ${ADMIN_PIN_DEF}
        echo 1 # RSA
        echo ${RSA_KEY_LENGTH} #Authentication key size set to RSA_KEY_LENGTH
        echo ${ADMIN_PIN_DEF}
    } | 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 "Setting key attributed to RSA ${RSA_KEY_LENGTH} bits in USB security dongle failed."
    fi
    # Generate OEM GPG keys
    {
        echo admin
        echo generate
        echo n
        echo ${ADMIN_PIN_DEF}
        echo ${USER_PIN_DEF}
        echo 0
        echo y
        echo ${GPG_USER_NAME} 
        echo ${GPG_USER_MAIL}
        echo ${GPG_USER_COMMENT}
    } | 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 | fold -s`
        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
    if [ "$CONFIG_TPM" = "y" ]; then
      tpm counter_create \
          -pwdo "$TPM_PASS" \
          -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"
    else
      ## needs to exist for initial call to unseal-hotp
      echo "0" > /boot/kexec_hotp_counter
    fi 

    # set default boot option
    set_default_boot_option

    # generate hashes
    find /boot -type f ! -name '*kexec*' -print0 \
        | xargs -0 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" \
            --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 | fold -s)
        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
    # 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 "/boot/loader/entries" ]; then
      for i in `find /boot -name "grub.cfg"`; do
        kexec-parse-bls "/boot" "$i" "/boot/loader/entries" >> $option_file
      done
    fi
    [ ! -s $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

    # get correct index for entry
    index=$(grep -n "$entry" $option_file | cut -f1 -d ':')

    # write new config
    echo "$entry" > /boot/kexec_default.$index.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 / Re-Ownership"
fi
if [ "$2" != "" ]; then
    bg_color=$2
else
    bg_color=""
fi

# show warning prompt
if [ "$CONFIG_TPM" = "y" ]; then
    TPM_STR="          * ERASE the TPM and own it with a password\n"
else
    TPM_STR=""
fi
if ! whiptail --yesno "
        This operation will automatically:\n
$TPM_STR
          * ERASE any keys or passwords on the GPG smart card,\n
            reset it to a factory state, generate new keys\n
            and optionally set custom PIN(s)
          * 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?" \
        $HEIGHT $WIDTH $CONTINUE $CANCEL $CLEAR $bg_color --title "$title_text" ; then
    exit 1
fi

# Inform user of security components affected for the following prompts
echo -e "The following security components will be provisioned with defaults or chosen PINs/passwords:
$CUSTOM_PASS_AFFECTED_COMPONENTS\n"

# Prompt to change default passwords
echo -e -n "Would you like to set a single custom password that will be provisioned to all security components? [y/N]: "
read -n 1 prompt_output
echo
if [ "$prompt_output" == "y" \
  -o "$prompt_output" == "Y" ] \
; then
  echo -e "\nThe chosen custom password must be at least 8 characters in length.\n"
  echo
  while [[  ${#CUSTOM_SINGLE_PASS} -lt 8 ]] ; do
    echo -e -n "Enter the custom password: "
    read CUSTOM_SINGLE_PASS
  done
  echo
  TPM_PASS=$CUSTOM_SINGLE_PASS
  USER_PIN=$CUSTOM_SINGLE_PASS
  ADMIN_PIN=$CUSTOM_SINGLE_PASS
else
  echo -e -n "Would you like to set distinct PINs/passwords to be provisioned to security components? [y/N]: "
  read -n 1 prompt_output
  echo
  if [ "$prompt_output" == "y" \
    -o "$prompt_output" == "Y" ] \
  ; then
    echo -e "\nThey must be each at least 8 characters in length.\n"
    echo
    if [ "$CONFIG_TPM" = "y" ]; then
      while [[  ${#TPM_PASS} -lt 8 ]] ; do
        echo -e -n "Enter desired TPM Ownership password: "
        read TPM_PASS
      done
    fi
    while [[  ${#ADMIN_PIN} -lt 8 ]] ; do
      echo -e -n "Enter desired GPG Admin PIN: "
      read ADMIN_PIN
    done
    while [[  ${#USER_PIN} -lt 8 ]] ; do
      echo -e -n "Enter desired GPG User PIN: "
      read USER_PIN
    done
    echo
  fi
fi

# If nothing is stored in custom variables, we set them to their defaults
if [ "$TPM_PASS" == "" ]; then TPM_PASS=$TPM_PASS_DEF; fi
if [ "$USER_PIN" == "" ]; then USER_PIN=$USER_PIN_DEF; fi
if [ "$ADMIN_PIN" == "" ]; then ADMIN_PIN=$ADMIN_PIN_DEF; fi

# Prompt to change default GnuPG key information
echo -e -n "Would you like to set custom user information for the GnuPG key? [y/N]: "
read -n 1 prompt_output
echo
if [ "$prompt_output" == "y" \
  -o "$prompt_output" == "Y" ] \
; then
	echo -e "\n\n"
	echo -e "We will generate a GnuPG (PGP) keypair identifiable with the following text form:"
	echo -e "Real Name (Comment) email@address.org"

	echo -e "\nEnter your Real Name (At least 5 characters long):"
	read -r GPG_USER_NAME
	while [[ ${#GPG_USER_NAME} -lt 5 ]]; do
	{
		echo -e "\nEnter your Real Name (At least 5 characters long):"
		read -r GPG_USER_NAME
	};done

	echo -e "\nEnter your email@adress.org:"
	read -r GPG_USER_MAIL
	while ! $(expr "$GPG_USER_MAIL" : '.*@' >/dev/null); do
	{
		echo -e "\nEnter your email@address.org:"
		read -r GPG_USER_MAIL
	};done

	echo -e "\nEnter Comment (Optional, to distinguish this key from others with same previous attributes. Must be smaller then 60 characters):"
	read -r GPG_USER_COMMENT
	while [[ ${#GPG_USER_COMMENT} -gt 60 ]]; do
	{
		echo -e "\nEnter Comment (Optional, to distinguish this key from others with same previous attributes. Must be smaller then 60 characters):"
		read -r GPG_USER_COMMENT
	};done
fi

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

# Prompt to insert USB drive if desired
echo -e -n "Would you like to export your public key to an USB drive? [y/N]: "
read -n 1 prompt_output
echo
if [ "$prompt_output" == "y" \
  -o "$prompt_output" == "Y" ] \
; then
    GPG_EXPORT=1
    # mount USB over /media only if not already mounted
    if ! grep -q /media /proc/mounts ; then
      # mount USB in rw
      if ! mount-usb rw 2>/tmp/error; then
        ERROR=$(tail -n 1 /tmp/error | fold -s)
        whiptail_error_die "Unable to mount USB on /media:\n\n${ERROR}"
      fi
    else
      #/media already mounted, make sure it is in r+w mode
      if ! mount -o remount,rw /media 2>/tmp/error; then
        ERROR=$(tail -n 1 /tmp/error | fold -s)
        whiptail_error_die "Unable to remount in read+write USB on /media:\n\n${ERROR}"
      fi
    fi
else
    GPG_EXPORT=0
    # needed for USB Security dongle below and is ensured via mount-usb in case of GPG_EXPORT=1
    enable_usb
fi

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

# detect and set /boot device
echo -e "\nDetecting and setting boot device...\n"
if ! detect_boot_device ; then
  SKIP_BOOT="y"
else
  echo -e "Boot device set to $CONFIG_BOOT_DEV\n"
fi

# update configs
if [[ "$SKIP_BOOT" == "n" ]]; then
  replace_config /etc/config.user "CONFIG_BOOT_DEV" "$CONFIG_BOOT_DEV"
  combine_configs
fi

## reset TPM and set password
if [ "$CONFIG_TPM" = "y" ]; then
  echo -e "\nResetting TPM...\n"
  {
      echo $TPM_PASS
      echo $TPM_PASS
  } | /bin/tpm-reset >/dev/null 2>/tmp/error
  if [ $? -ne 0 ]; then
      ERROR=$(tail -n 1 /tmp/error | fold -s)
      whiptail_error_die "Error resetting TPM:\n\n${ERROR}"
  fi
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 around 3 minutes...)\n"
gpg_key_reset

# 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"

#Applying custom GPG PINs
if [ "$USER_PIN" != "" -o "$ADMIN_PIN" != "" ]; then
  echo -e "\nChanging default GPG Admin PIN\n"
  gpg_key_change_pin "3" "$ADMIN_PIN_DEF" "$ADMIN_PIN"
  echo -e "\nChanging default GPG User PIN\n"
  gpg_key_change_pin "1" "$USER_PIN_DEF" "$USER_PIN"
fi

# export pubkey to file
if ! gpg --export --armor $GPG_GEN_KEY > "${PUBKEY}" 2>/tmp/error ; then
    ERROR=$(tail -n 1 /tmp/error | fold -s)
    whiptail_error_die "GPG Key gpg export to file failed!\n\n$ERROR"
fi

## export pubkey to USB
if [ $GPG_EXPORT -ne 0 ]; then
    echo -e "\nExporting generated key to USB...\n"
    # copy to USB
    if ! cp "${PUBKEY}" "/media/${GPG_GEN_KEY}.asc" 2>/tmp/error ; then
        ERROR=$(tail -n 1 /tmp/error | fold -s)
        whiptail_error_die "Key export error: unable to copy ${GPG_GEN_KEY}.asc to /media:\n\n$ERROR"
    fi
    mount -o remount,ro /media 2>/dev/null
fi

## 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 | fold -s)
    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 | fold -s)
    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 | fold -s)
    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 | fold -s)
    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 | fold -s)
    whiptail_error_die "Error flashing updated firmware image:\n\n$ERROR"
fi

## sign files in /boot and generate checksums
if [[ "$SKIP_BOOT" == "n" ]]; then
  echo -e "\nSigning boot files and generating checksums...\n"
  generate_checksums
fi

## Show user current provisioned PINS/Password prior of reboot
whiptail --msgbox "
    TPM Owner Password: $TPM_PASS\n
    GPG Admin PIN: $ADMIN_PIN\n
    GPG User PIN: $USER_PIN\n\n" \
    $HEIGHT $WIDTH --title "Provisioned secrets"

## all done -- reboot
whiptail --msgbox "
    OEM Factory Reset / Re-Ownership 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 Enter to reboot.\n" \
    $HEIGHT $WIDTH --title "OEM Factory Reset / Re-Ownership Complete"

reboot