#!/bin/bash

##==============================================================================
##
## This script prepares a platform for LSVM booting.
##
##==============================================================================

##==============================================================================
##
## Load utility functions
##
##==============================================================================
. ./scripts/blkdev_utils

##==============================================================================
##
## Process command-line options:
##
##==============================================================================

for opt
do

  arg=`expr "x$opt" : 'x[^=]*=\(.*\)'`

  case $opt in

    -h | --help)
      help=1
      ;;

    --secureboot)
      securebootopt=1
      ;;

    --seal)
      sealopt=1
      ;;

    --dbxupdate)
      dbxupdateopt=1
      ;;

    *)
      echo "$0: unknown option:  $opt"
      exit 1
      ;;

  esac

done

##==============================================================================
##
## Apply environment variable options:
##
##==============================================================================

if [ "${LSVMPREP_SEAL}" == "1" ]; then
    sealopt=1
fi

if [ "${LSVMPREP_DBXUPDATE}" == "1" ]; then
    dbxupdateopt=1
fi

##==============================================================================
##
## help:
##
##==============================================================================

if [ "$help" = "1" ]; then
    cat<<EOF

OVERVIEW:

Prepares a Linux VM to be booted as a Linux shielded VM.

    $ ./lsvmprep [OPTIONS]

OPTIONS:
    -h, --help              Print this help message.
    --seal                  Seal and create /boot/efi/EFI/boot/sealedkeys with
                            well-known passphrases (i.e., 'passphrase').
    --dbxupdate             Update the UEFI dbx variable.
    --secureboot            Seal for secure boot (if --seal option present).

EOF
    exit 0
fi

##==============================================================================
##
## chkerr()
##
##==============================================================================

chkerr()
{
    local rc;
    local msg;

    if [ "$#" != "2" ]; then
        echo "$0: $FUNCNAME(): wrong # of args"
        exit 1
    fi

    rc=$1
    msg=$2

    if [ "$rc" == "0" ]; then
        return 0
    fi

    echo "$0: $msg"
    exit 1
}


##==============================================================================
##
## Define top variable:
##
##==============================================================================

top=`pwd`

##==============================================================================
##
## Resolve EFI vendor directory name:
##
##==============================================================================

vendor=`grep "^ID=" /etc/os-release | cut -d "=" -f2 | sed "s/\"//g"`

case ${vendor} in
    centos)
        ;;
    rhel)
        vendor="redhat"
        ;;
    sles)
        ;;
    ubuntu)
        ;;
    *)
        echo "$0: unsupported vendor distribution: ${vendor}"
        exit 1
esac

##==============================================================================
##
## Ask for permissio to change the system
##
##==============================================================================


ask_permission()
{
cat<<EOF

***************************************************
*     ____    _   _   _ _____ ___ ___  _   _      *
*    / ___|  / \ | | | |_   _|_ _/ _ \| \ | |     *
*   | |     / _ \| | | | | |  | | | | |  \| |     *
*   | |___ / ___ \ |_| | | |  | | |_| | |\  |     *
*    \____/_/   \_\___/  |_| |___\___/|_| \_|     *
*           					  *
*           					  *
* LSVMPREP is about to encrypt the boot partition *
* and make irreversible configuration changes to  *
* this machine. If you are certain you want to    *
* proceed, type YES in uppercase and then press   *
* enter; else press ENTER to terminate.           *
*          					  *
***************************************************

EOF

    echo -n "> "
    read answer
    echo ""

    if [ "${answer}" != "YES" ]; then
        echo "$0: no action taken"
        echo ""
        exit 1
    fi
}

ask_permission

##==============================================================================
##
## check_root()
##
##==============================================================================

check_root_passphrase()
{
    rootdev=`./scripts/rootdev`

    if [ -z "${rootdev}" ]; then
        echo "$0: failed to resolve root partition"
        exit 1
    fi

    # Check whether cryptsetup program exists.
    if [ ! -x "/sbin/cryptsetup" ]; then
        echo "$0: ERROR: /sbin/cryptsetup not found"
        echo ""
        exit 1
    fi

    # Check whether /lib/cryptsetup/askpass program exists.
    if [ "${vendor}" == "ubuntu" ]; then
        if [ ! -x "/lib/cryptsetup/askpass" ]; then
            echo "$0: ERROR: /lib/cryptsetup/askpass dependency not found"
            echo ""
            exit 1
        fi
    fi

    # Check whether root partition is encrypted.
    /sbin/cryptsetup luksDump $rootdev 2> /dev/null > /dev/null
    if [ "$?" != "0" ]; then
        echo "$0: ERROR: The root partition is not a LUKS partition: $rootdev"
        echo ""
        exit 1
    fi

    # Check whether root partition is encrypted with "passprhase"
    echo -n "passphrase" | /sbin/cryptsetup luksDump --dump-master-key --key-file=- $rootdev 2> /dev/null > /dev/null
    if [ "$?" != "0" ]; then
        echo "$0: ERROR: The root partition passphrase must be 'passphrase'"
        echo ""
        exit 1
    fi
}

check_root_passphrase

##==============================================================================
##
## Resolve 'shim':
##
##==============================================================================

shim=`ls /boot/efi/EFI/${vendor}/shim*.efi | grep -E "shim(x64)?.efi"`

if [ ! -x "${shim}" ]; then
    echo "$0: shim not found"
    exit 1
fi

##==============================================================================
##
## Resolve 'grub':
##
##==============================================================================

grub=`ls /boot/efi/EFI/${vendor}/grub*.efi | grep -E "grub(x64)?.efi"`
if [ "$vendor" = "sles" ]; then
    grub=`ls /boot/efi/EFI/sles/grub.efi`
fi

if [ ! -x "${grub}" ]; then
    echo "$0: grub not found"
    exit 1
fi

##==============================================================================
##
## Resolve location of lsvmtool.
##
##==============================================================================

lsvmtool=${top}/build/bin/lsvmtool

if [ ! -x "${lsvmtool}" ]; then
    echo "$0: lsvmtool not found: ${lsvmtool}"
    exit 1
fi

##==============================================================================
##
## check_for_revocation() -- check whether lsvmload.efi is revoked.
##
##==============================================================================

check_for_revocation()
{
    local lsvmload=${top}/lsvmload/lsvmload.efi
    local dbxupdate=${top}/dbxupdate.bin

    if [ "$#" != "0" ]; then
        echo "$0: $FUNCNAME(): wrong # of args"
        exit 1
    fi

    echo "Checking ${lsvmload} for revocation"

    result=`${lsvmtool} revoked ${lsvmload}`
    chkerr "$?" "$0: lsvmtool revoked check failed"
    if [ "${result}" == "yes" ]; then
        echo "$0: ${lsvmtool} found on UEFI DBX revocation list"
        exit 1
    fi

    if [ -f ${dbxupdate} ]; then
        result=`${lsvmtool} revoked ${lsvmload} ${dbxupdate}`
        chkerr "$?" "$0: lsvmtool revoked check failed"
        if [ "${result}" == "yes" ]; then
            echo "$0: ${lsvmtool} found in ${dbxupdate}"
            exit 1
        fi
    fi

    chkerr "$?" "$0: failed to install lsvmload"
}

if [ "${dbxupdateopt}" == "1" ]; then
    check_for_revocation
fi

##==============================================================================
##
## Remove /boot/efi/EFI/boot/lsvmprep
##
##==============================================================================

lsvmprepfile=/boot/efi/EFI/boot/lsvmprep
rm -f ${lsvmprepfile}

##==============================================================================
##
## Encrypt the boot drive:
##
##==============================================================================

${top}/scripts/encryptboot
chkerr "$?" "$0: failed to encrypt the boot drive"

##==============================================================================
##
## Set up the root drive:
##
##==============================================================================
if [[ ${vendor} != "ubuntu" ]]; then
    ${top}/scripts/setuproot
    chkerr "$?" "$0: failed to setup root drive"
fi

##==============================================================================
##
## install_lsvmtool(lsvmtool)
##
##==============================================================================

install_lsvmtool()
{
    if [ "$#" != "1" ]; then
        echo "$0: $FUNCNAME(): wrong # of args"
        exit 1
    fi

    local lsvmtool=$1

    if [ ! -x "${lsvmtool}" ]; then
        echo "$0: $FUNCNAME(): missing executable: ${lsvmtool}"
        exit 1
    fi

    local filename=/sbin/lsvmtool

    echo "Creating ${filename}"
    cp ${lsvmtool} ${filename}

    if [ ! -x "${filename}" ]; then
        echo "$0: $FUNCNAME(): failed to create: ${filename}"
        exit 1
    fi
}

install_lsvmtool ${lsvmtool}

##==============================================================================
##
## install_lsvmload()
##
##==============================================================================

install_lsvmload()
{
    local lsvmload=${top}/lsvmload/lsvmload.efi
    local bootx64=/boot/efi/EFI/boot/bootx64.efi

    if [ "$#" != "0" ]; then
        echo "$0: $FUNCNAME(): wrong # of args"
        exit 1
    fi

    echo "Creating ${bootx64}"

    if [ ! -x "${lsvmload}" ]; then
        echo "$0: not found: lsvmload/lsvmload.efi"
        exit 1
    fi

    install -D ${lsvmload} ${bootx64}
    chkerr "$?" "$0: failed to install lsvmload"
}

install_lsvmload

##==============================================================================
##
## install_lsvmconf(vendor, shim, grub)
##
##==============================================================================

install_lsvmconf()
{
    local vendor=$1
    local shim=$2
    local grub=$3

    local dirname=/boot/efi/EFI/boot
    local filename=${dirname}/lsvmconf
    local tempfile=`/bin/mktemp`

    if [ "$#" != "3" ]; then
        echo "$0: $FUNCNAME(): wrong # of args"
        exit 1
    fi

    if [ -z "${vendor}" ]; then
        echo "$0: $FUNCNAME(): vendor undefined"
        exit 1
    fi

    if [ ! -x "${shim}" ]; then
        echo "$0: $FUNCNAME(): file not found: $shim"
        exit 1
    fi

    if [ ! -x "${grub}" ]; then
        echo "$0: $FUNCNAME(): file not found: $grub"
        exit 1
    fi

    local bootdev=`./scripts/bootdev`
    local bootdevbase=`basename $bootdev`

    if [ -z "${bootdev}" ]; then
        echo "$0: $FUNCNAME(): failed to resolve boot device"
        exit 1
    fi

    local rootdev=`./scripts/rootdev`
    local rootdevbase=`basename $rootdev`

    if [ -z "${rootdev}" ]; then
        echo "$0: $FUNCNAME(): failed to resolve root device"
        exit 1
    fi

    rm -rf ${tempfile}

    echo "EFIVendorDir=${vendor}" >> ${tempfile}
    echo "BootDeviceLUKS=`${lsvmtool} partuuid ${bootdev}`" >> ${tempfile}
    echo "RootDeviceLUKS=`${lsvmtool} partuuid ${rootdev}`" >> ${tempfile}
    echo "BootDevice=$(ls -l /dev/disk/by-partuuid | grep "$bootdevbase" | awk '{print $(NF-2)}')" >> ${tempfile}
    echo "RootDevice=$(ls -l /dev/disk/by-partuuid | grep "$rootdevbase" | awk '{print $(NF-2)}')" >> ${tempfile}

    echo "Creating ${filename}"
    install -D ${tempfile} ${filename}
    chkerr "$?" "$0: $FUNCNAME(): failed to install ${filename}"

    rm -rf ${tempfile}
}

install_lsvmconf ${vendor} ${shim} ${grub}

##==============================================================================
##
## install_shim()
##
##==============================================================================

install_shim()
{
    local shim=$1

    if [ "$#" != "1" ]; then
        echo "$0: $FUNCNAME(): wrong # of args"
        exit 1
    fi

    if [ ! -x "${shim}" ]; then
        echo "$0: $FUNCNAME(): executable file not found: ${shim}"
        exit 1
    fi

    local dirname=/boot/lsvmload
    local filename="shimx64.efi"

    echo "Creating ${dirname}/${filename}"

    mkdir -p ${dirname}

    if [ ! -d "${dirname}" ]; then
        echo "$0: $FUNCNAME(): directory not found: ${dirname}"
        exit 1
    fi

    rm -rf ${dirname}/${filename}

    cp ${shim} ${dirname}/${filename}
    chkerr "$?" "$0: $FUNCNAME(): failed to install ${shim}"
}

install_shim ${shim}

##==============================================================================
##
## install_grub()
##
##==============================================================================

install_grub()
{
    local grub=$1

    if [ "$#" != "1" ]; then
        echo "$0: $FUNCNAME(): wrong # of args"
        exit 1
    fi

    if [ ! -x "${grub}" ]; then
        echo "$0: $FUNCNAME(): executable file not found: ${grub}"
        exit 1
    fi

    local dirname=/boot/lsvmload
    local filename="grubx64.efi"

    echo "Creating ${dirname}/${filename}"

    mkdir -p ${dirname}

    if [ ! -d "${dirname}" ]; then
        echo "$0: $FUNCNAME(): directory not found: ${dirname}"
        exit 1
    fi

    rm -rf ${dirname}/${filename}

    cp ${grub} ${dirname}/${filename}
    chkerr "$?" "$0: $FUNCNAME(): failed to install ${grub}"
}

install_grub ${grub}

##==============================================================================
##
## patch_grub_timeout()
##
##     GRUB_TIMEOUT=0
##     GRUB_HIDDEN_TIMEOUT=0
##     GRUB_RECORDFAIL_TIMEOUT=0
##
##==============================================================================

function patch_grub_timeout()
{
    local filename=/etc/default/grub

    # Fail if /etc/default/grub does not exist:
    if [ ! -f "${filename}" ]; then
        return 1;
    fi

    # Perform substitutions:
    local sed1='s/^[ \t]*#\?[ \t]*GRUB_TIMEOUT[ \t]*=.*/GRUB_TIMEOUT=0/g'
    local sed2='s/^[ \t]*#\?[ \t]*GRUB_HIDDEN_TIMEOUT[ \t]*=.*/GRUB_HIDDEN_TIMEOUT=0/g'
    local sed3='s/^[ \t]*#\?[ \t]*GRUB_RECORDFAIL_TIMEOUT[ \t]*=.*/GRUB_RECORDFAIL_TIMEOUT=0/g'
    local tmpfile=`/bin/mktemp`
    cat ${filename} | sed "${sed1}" | sed "${sed2}" | sed "${sed3}" > ${tmpfile}

    # Verify that file contains single "GRUB_TIMEOUT=0":
    match=`grep '^GRUB_TIMEOUT=0$' ${tmpfile}`
    if [ "${match}" != "GRUB_TIMEOUT=0" ]; then
        echo "$0: $FUNCNAME(): patch failed (1): ${filename}"
        exit 1
    fi

    # If file does not contain "GRUB_HIDDEN_TIMEOUT=0" then force it:
    match=`grep '^GRUB_HIDDEN_TIMEOUT=0$' ${tmpfile}`
    if [ "$?" != "0" ]; then
        echo "GRUB_HIDDEN_TIMEOUT=0" >> ${tmpfile}
    fi

    # Verify that file contains single "GRUB_HIDDEN_TIMEOUT=0":
    match=`grep '^GRUB_HIDDEN_TIMEOUT=0$' ${tmpfile}`
    if [ "${match}" != "GRUB_HIDDEN_TIMEOUT=0" ]; then
        echo "$0: $FUNCNAME(): patch failed (2): ${filename}"
        exit 1
    fi

    # If file does not contain "GRUB_RECORDFAIL_TIMEOUT=0" then force it:
    match=`grep '^GRUB_RECORDFAIL_TIMEOUT=0$' ${tmpfile}`
    if [ "$?" != "0" ]; then
        echo "GRUB_RECORDFAIL_TIMEOUT=0" >> ${tmpfile}
    fi

    # Verify that file contains one "GRUB_RECORDFAIL_TIMEOUT=0":
    match=`grep '^GRUB_RECORDFAIL_TIMEOUT=0$' ${tmpfile}`
    if [ "${match}" != "GRUB_RECORDFAIL_TIMEOUT=0" ]; then
        echo "$0: $FUNCNAME(): patch failed (3): ${filename}"
        exit 1
    fi

    # Replace /etc/default/grub:
    cp ${tmpfile} ${filename}
    chkerr "$?" "$0: $FUNCNAME(): patch failed (4): ${filename}"

    rm -f ${tmpfile}
    return 0
}

patch_grub_timeout

##==============================================================================
##
## patch_grub_terminal()
##
##     GRUB_TERMINAL=console
##
##==============================================================================

function patch_grub_terminal()
{
    local filename=/etc/default/grub

    # Fail if /etc/default/grub does not exist:
    if [ ! -f "${filename}" ]; then
        return 1;
    fi

    # Perform substitutions:
    local sed1='s/^[ \t]*#\?[ \t]*GRUB_TERMINAL[ \t]*=.*/GRUB_TERMINAL=console/g'
    local tmpfile=`/bin/mktemp`
    cat ${filename} | sed "${sed1}" > ${tmpfile}

    # If file does not contain single "GRUB_TERMINAL=console", then append:
    match=`grep '^GRUB_TERMINAL=console$' ${tmpfile}`
    if [ "${match}" != "GRUB_TERMINAL=console" ]; then
        echo "GRUB_TERMINAL=console" >> ${tmpfile}
    fi

    # Replace /etc/default/grub:
    cp ${tmpfile} ${filename}
    chkerr "$?" "$0: $FUNCNAME(): patch failed (5): ${filename}"

    rm -f ${tmpfile}

    return 0
}

patch_grub_terminal

##==============================================================================
##
## generate_grubcfg(top, vendor, lsvmtool, grub)
##
##==============================================================================

generate_grubcfg()
{
    local top=$1
    local vendor=$2
    local lsvmtool=$3
    local grub=$4

    if [ "$#" != "4" ]; then
        echo "$0: $FUNCNAME(): wrong # of args"
        exit 1
    fi

    if [ -z "${vendor}" ]; then
        echo "$0: $FUNCNAME(): vendor undefined"
        exit 1
    fi

    if [ ! -x "${lsvmtool}" ]; then
        echo "$0: $FUNCNAME(): file not found: $lsvmtool"
        exit 1
    fi

    if [ ! -x "${grub}" ]; then
        echo "$0: $FUNCNAME(): file not found: $grub"
        exit 1
    fi

    case ${vendor} in
        centos|redhat)
            # Relocate 'grub.cfg' from ESP to boot partition:
            rm -f /etc/grub2-efi.cfg
            ln -s /boot/grub2/grub.cfg /etc/grub2-efi.cfg
            chkerr "$?" "$0: $FUNCNAME(): failed to link /etc/grub2-efi.cfg"
            grubby --update-kernel=ALL --config-file=/etc/grub2-efi.cfg

            local grubcfg=/boot/grub2/grub.cfg
            local mkconfig=grub2-mkconfig
            ${top}/scripts/${vendor}/update_grub_config
            chkerr "$?" "$0: $FUNCNAME(): update_grub_config failed"
            ;;

       ubuntu)
            local grubcfg=/boot/grub/grub.cfg
            local mkconfig=grub-mkconfig
            ;;

       sles)
            local grubcfg=/boot/grub2/grub.cfg
            local mkconfig=grub2-mkconfig
            ;;

    esac

    echo "Creating ${grubcfg}"
    ${mkconfig} > ${grubcfg} 2> /dev/null
    chkerr "$?" "$0: $FUNCNAME(): ${grubmkconfig} failed"
}

generate_grubcfg ${top} ${vendor} ${lsvmtool} ${grub}

##==============================================================================
##
## configure_initrd(top, vendor)
##
##==============================================================================

configure_initrd()
{
    local top=$1
    local vendor=$2
    local dirname=${top}/scripts/${vendor}

    if [ ! -d "${dirname}" ]; then
        echo "$0: $FUNCNAME(): no such directory: ${dirname}"
        exit 1
    fi

    if [ ! -x "${dirname}/patch_initrd" ]; then
        echo "$0: $FUNCNAME(): file not found: ${dirname}/patch_initrd"
        exit 1
    fi

    cd ${dirname}
    chkerr "$?" "$0: $FUNCNAME() failed to change to directory: ${dirname}"

    ./patch_initrd
    chkerr "$?" "$0: $FUNCNAME() failed to execute ${dirname}/patch_initrd"

    cd ${top}
}

configure_initrd ${top} ${vendor}

##==============================================================================
##
## generate_initrd(vendor)
##
##==============================================================================

generate_initrd()
{
    local vendor=$1
    local initrd=$2

    echo "Updating initial ramdisks"

    case ${vendor} in
        centos|redhat)
            dracut -f --regenerate-all
            chkerr "$?" "$0: $FUNCNAME(): failed to update initrd"
            ;;
        sles)
            dracut -f --regenerate-all
            chkerr "$?" "$0: $FUNCNAME(): failed to update initrd"
            ;;
        ubuntu)
            update-initramfs -u -k all 2> /dev/null > /dev/null
            chkerr "$?" "$0: $FUNCNAME(): failed to update initrd"
            ;;
    esac
}

generate_initrd ${vendor}

##==============================================================================
##
## apply_dbxupdate()
##
##==============================================================================

apply_dbxupdate()
{
    local src="${top}/dbxupdate.bin"
    local dest="/boot/lsvmload/dbxupdate.bin"

    if [ -f "${src}" ]; then

        need=`${lsvmtool} needdbxupdate ${src}`
        chkerr "$?" "$0: $FUNCNAME(): failed to apply update"

        echo -n "Checking for DBX update: "

        if [ "${need}" == "yes" ]; then
            echo "created ${dest}"
            cp ${src} ${dest}
            chkerr "$?" "$0: $FUNCNAME(): failed to create ${dest}"
        else
            echo "up to date"
        fi
    fi
}

if [ "${dbxupdateopt}" == "1" ]; then
    apply_dbxupdate
fi

##==============================================================================
##
## prepare_keys()
##
##==============================================================================

prepare_keys()
{
    local lsvmtool=$1
    local bootkey=$2
    local rootkey=$3

    local tempboot=`/bin/mktemp`
    local temproot=`/bin/mktemp`
    local tempout=`/bin/mktemp`

    echo -n ${bootkey} > ${tempboot}
    echo -n ${rootkey} > ${temproot}


    ${lsvmtool} serializekeys ${tempboot} ${temproot} ${tempout}
    rm -rf ${tempboot}
    rm -rf ${temproot}
    echo -n ${tempout}
}


##==============================================================================
##
## seal_keys(top)
##
##==============================================================================

seal_key()
{
    local top=$1
    local lsvmtool=$2
    local keys=$3
    local filename=$4

    if [ "$#" != "4" ]; then
        echo "$0: $FUNCNAME(): wrong # of args"
        exit 1
    fi

    echo "Creating ${filename}"

    local lsvmload=${top}/lsvmload/lsvmload.efi

    if [ "${securebootopt}" == "1" ]; then
        ${lsvmtool} seallsvmloadpolicy --secureboot ${lsvmload} ${keys} ${filename}
    else
        ${lsvmtool} seallsvmloadpolicy ${lsvmload} ${keys} ${filename}
    fi

    if [ "$?" != "0" ]; then
        echo "$0: $FUNCNAME(): failed to seal key"
        exit 1
    fi
}

if [ "${sealopt}" == "1" ]; then

    bootkey=passphrase
    rootkey=passphrase
    combined_keys=$(prepare_keys ${lsvmtool} ${bootkey} ${rootkey})

    seal_key ${top} ${lsvmtool} ${combined_keys} /boot/efi/EFI/boot/sealedkeys

    rm -rf ${combined_keys}

fi

##==============================================================================
##
## Install VERSION file (as /boot/efi/EFI/boot/lsvmprep)
##
##==============================================================================

cp VERSION ${lsvmprepfile}

# Run sanity check:
./sanity

##==============================================================================
##
## final_word()
##
##==============================================================================

final_word()
{
cat<<EOF

This image has been successfully prepared for templatization.

The passphrase for the root partition is: passphrase
The passphrase for the boot partition is: passphrase

EOF
}

final_word
