#!/usr/bin/bash
set -euo pipefail
shopt -s expand_aliases

# Anti Evil Maid for dracut by Invisible Things Lab
# Copyright (C) 2010 Joanna Rutkowska <joanna@invisiblethingslab.com>
#
# Mount our device, read the sealed secret blobs, initialize TPM
# and finally try to unseal the secrets and display them to the user


MNT=/anti-evil-maid
UNSEALED_SECRET=/tmp/unsealed-secret
LUKS_HEADER_DUMP=/tmp/luks-header-dump
LUKS_PCR=13

PLYMOUTH_MESSAGES=()

plymouth_message() {
    if [ "${#PLYMOUTH_MESSAGES[@]}" -eq 0 ]; then
        # add vertical "padding" to avoid printing messages over plymouth's
        # prompt help
        plymouth message --text=""
        plymouth message --text=""
        plymouth message --text=""
    fi

    plymouth message --text="$*"
    PLYMOUTH_MESSAGES+=("$*")
}

plymouth_messages_hide() {
    for m in "${PLYMOUTH_MESSAGES[@]}"; do
        plymouth hide-message --text="$m"
    done
}

# shellcheck source=../sbin/anti-evil-maid-lib
. anti-evil-maid-lib


# find AEM device

UUID=$(getparams aem.uuid)
DEV=/dev/disk/by-uuid/$UUID

udevadm trigger
udevadm settle
waitfor -b "$DEV"

n=$(lsblk -nr -o UUID | grep -Fixc "$UUID") || true
if [ "$n" != 1 ]; then
    message "Error: found ${n:-?} devices with UUID $UUID"
    exit 1
fi

LABEL=$(lsblk -dnr -o LABEL "$DEV")
if [[ "$LABEL" != "$LABEL_PREFIX"* ]]; then
    message "AEM boot device $DEV has wrong label: $LABEL"
    exit 1
fi


# mount AEM device

log "Mounting $DEV (\"$LABEL\")..."
mkdir -p "$MNT"
mount -t ext4 -o ro "$DEV" "$MNT"
# this way an error prior to unmounting will keep the device usable after initrd
trap 'umount "$MNT"' EXIT


# setup TPM & copy secrets to initrd tmpfs

log "Initializing TPM..."
modprobe tpm_tis
validatetpm || exit 1
ip link set dev lo up
mkdir -p "${TPMS_DIR%/*}"
log "Copying sealed AEM secrets..."
cp -Tr "$MNT/aem/${TPMS_DIR##*/}" "${TPMS_DIR}"
tpmstartinitrdservices

SEALED_SECRET_TXT=$TPM_DIR/$LABEL/secret.txt.sealed2
SEALED_SECRET_KEY=$TPM_DIR/$LABEL/secret.key.sealed2
SEALED_SECRET_OTP=$TPM_DIR/$LABEL/secret.otp.sealed2
SEALED_SECRET_FRE=$TPM_DIR/$LABEL/secret.fre.sealed2


# unmount AEM device

log "Unmounting $DEV (\"$LABEL\")..."
umount "$MNT"
# remove umount trap set after mount
trap - EXIT

if [ "$(blockdev --getro "$DEV")" = 1 ]; then
    message "You should now unplug the AEM device if it is intentionally read-only."
fi


# Extend PCR with LUKS header(s)

getluksuuids |
sort -u |
while read -r luksid; do
    waitfor -b "/dev/disk/by-uuid/$luksid"

    cryptsetup luksHeaderBackup "/dev/disk/by-uuid/$luksid" \
               --header-backup-file "$LUKS_HEADER_DUMP"
    luks_header_hash=$(hashfile "$LUKS_HEADER_DUMP")
    log "Extending PCR $LUKS_PCR, value $luks_header_hash, device $luksid..."
    tpmpcrextend "$LUKS_PCR" "$luks_header_hash"
done


# cache suffix and SRK password, if applicable

mkdir -p "$CACHE_DIR"
echo "${LABEL##"$LABEL_PREFIX"}" >"$SUFFIX_CACHE"

Z=$(tpmzsrk)

if [ -n "$Z" ]; then
    true >"$SRK_PASSWORD_CACHE"
else
    for _ in 1 2 3; do
        log "Prompting for SRK password..."

        if systemd-ask-password --timeout=0 \
                                "TPM SRK password to unseal the secret(s)" \
                                > "$SRK_PASSWORD_CACHE" && checksrkpass; then
             log "Correct SRK password"
             break
        fi

        log "Wrong SRK password"
    done
fi


# check freshness token

log "Unsealing freshness token..."
if tpmunsealdata "$Z" "$SEALED_SECRET_FRE" "$UNSEALED_SECRET" \
                 "$TPM_DIR/$LABEL"; then
    log "Freshness token unsealed."
    true >"$CACHE_DIR/unseal-success"
else
    log "Freshness token unsealing failed!"
    log "This is expected during the first boot from a particular"
    log "AEM media or after updating any of the boot components or"
    log "changing their configuration."
    exit 1
fi

if checkfreshness "$UNSEALED_SECRET"; then
    log "Freshness token valid, continuing."
else
    log "Freshness token invalid!"
    exit 1
fi


# unseal & show OTP if provisioned
# unseal & decrypt key file unless the user switches to text secret mode

if [ -e "$SEALED_SECRET_OTP" ]; then
    alias otp=true
else
    alias otp=false
fi

if otp; then
    log "Unsealing TOTP shared secret seed..."
    if tpmunsealdata "$Z" "$SEALED_SECRET_OTP" "$UNSEALED_SECRET" \
                     "$TPM_DIR/$LABEL"; then
        log "TOTP secret unsealed."

        message ""
        message "Never type in your key file password unless the code below is correct!"
        message ""

        seed=$(cat "$UNSEALED_SECRET")
        last=
        {
            trap 'plymouth_messages_hide; exit' TERM
            while :; do
                now=$(date +%s)
                if [ -z "$last" ] ||
                   { [ "$((now % 30))" = 0 ] && [ "$last" != "$now" ]; }; then
                    code=$(oathtool --totp -b "$seed")
                    message "[ $(date) ] TOTP code: $code"
                    last=$now
                fi
                sleep 0.1
            done
        } &
        totp_loop_pid=$!

        if tpmunsealdata "$Z" "$SEALED_SECRET_KEY" "$UNSEALED_SECRET" \
                         "$TPM_DIR/$LABEL"; then
            for _ in 1 2 3; do
                pass=$(systemd-ask-password --timeout=0 \
                       'LUKS key file password (or "t" to show text secret)')

                if [ "$pass" = "t" ]; then
                    alias otp=false
                    break
                fi

                if scrypt dec -P "$UNSEALED_SECRET" /tmp/aem-keyfile \
                   <<<"$pass"; then
                    log "Correct LUKS key file password"
                    # dracut "90crypt" module will parse the
                    #   rd.luks.key=/tmp/aem-keyfile
                    # kernel cmdline arg and attempt to use it;
                    # this file is deleted on root switch
                    # along with everything in /tmp
                    break
                else
                    log "Wrong LUKS key file password"
                fi
            done
        fi

        kill "$totp_loop_pid"
    fi
fi


# unseal text secret

if ! otp; then
    log "Unsealing text secret..."
    if tpmunsealdata "$Z" "$SEALED_SECRET_TXT" "$UNSEALED_SECRET" \
                     "$TPM_DIR/$LABEL"; then
        {
            message ""
            message "$(cat "$UNSEALED_SECRET" 2>/dev/null)"
            message ""
        } 2>&1  # don't put the secret into the journal
        message "Never type in your disk password unless the secret above is correct!"
        waitforenter
    fi

    plymouth_messages_hide
    clear
fi


# prevent sealing service from starting if user unplugged
# the (supposedly read-only) AEM device

if [ ! -b "$DEV" ]; then
    rm -rf "$CACHE_DIR"
fi
