LABEL_PREFIX=aem
SYSFS_TPM_DIR=/sys/class/tpm/tpm0
AEM_DIR=/var/lib/anti-evil-maid
TPM_DIR=/var/lib/tpm
TPMS_DIR=${TPM_DIR}s
CACHE_DIR=/run/anti-evil-maid
SRK_PASSWORD_CACHE=$CACHE_DIR/srk-password
SUFFIX_CACHE=$CACHE_DIR/suffix
TPM_OWNER_PASSWORD_FILE=$AEM_DIR/tpm-owner-pw
TPM_FRESHNESS_PASSWORD_FILE=$AEM_DIR/tpm-freshness-pw
TPM_FRESHNESS_INDEX="0x454d"
TPM_FRESHNESS_SLOTS=8


# work with or without plymouth

if command plymouth --ping 2>/dev/null; then
    alias plymouth_active=true
    alias message=plymouth_message
else
    alias plymouth=:
    alias plymouth_active=false
    alias message=log
fi


getparams() {
    _CMDLINE=${_CMDLINE-$(cat /proc/cmdline)}

    for _param in $_CMDLINE; do
        for _key; do
            case "$_param" in "$_key"=*)
                printf '%s\n' "${_param#*=}"
                break
            esac
        done
    done
}

getluksuuids() {
    getparams rd.luks.uuid rd_LUKS_UUID | sed s/^luks-//
}

log() {
    echo "${0##*/}: $1" >&2
}

hex() {
    xxd -ps | tr -dc 0-9a-f
}

unhex() {
    tr -dc 0-9a-f | xxd -ps -r
}

waitfor() {
    case $# in
        2) _file=$2; _what=connected ;;
        3) _file=$3; _what=removed ;;
        *) return 1 ;;
    esac

    if [ "$@" ]; then
        return
    fi

    message "Waiting for $_file to be $_what..."
    plymouth pause-progress
    until [ "$@" ]; do
        sleep 0.1
    done
    plymouth unpause-progress
    message "$_file $_what"
}

waitforenter() {
    msg='Press <ENTER> to continue...'
    if plymouth_active; then
        message "$msg"
        plymouth watch-keystroke --keys=$'\n'
    else
        systemd-ask-password --timeout=0 "$msg" >/dev/null
    fi
}

checktpmnvram() {
    # checks whether the TPM NVRAM area is defined
    # NOTE: tpm_nvinfo does not return non-zero if requested index
    # is not a defined NVRAM area so we need to parse
    if ! tpm_nvinfo -i "$TPM_FRESHNESS_INDEX" | grep -q 'AUTHWRITE'; then
        return 1
    fi
}

createtpmnvram() {
    # create the world-readable/AUTHWRITE TPM NVRAM area to hold up to
    # TPM_FRESHNESS_SLOTS anti-replay freshness token hashes;
    # takes TPM owner password as an agument
    if [ ! -e "$TPM_FRESHNESS_PASSWORD_FILE" ]; then
        message "Generating TPM NVRAM area AUTHWRITE password"
        head -c 16 /dev/random | hex > "$TPM_FRESHNESS_PASSWORD_FILE"
    fi

    _pw=$(cat "$TPM_FRESHNESS_PASSWORD_FILE")

    if ! tpm_nvdefine -i "$TPM_FRESHNESS_INDEX" \
            -s $((TPM_FRESHNESS_SLOTS * 20)) \
            -p AUTHWRITE --pwda="$_pw" --pwdo="$1"; then
        return 1
    fi
}

suffixtoslotfile() {
    echo "$AEM_DIR/$LABEL_PREFIX$1/tpm-freshness-slot"
}

suffixtoslot() {
    # returns the slot number assigned to the AEM media given its label suffix
    # as the first argument
    _slotfile=$(suffixtoslotfile "$1")
    cat "$_slotfile" 2>/dev/null
}

assignslottosuffix() {
    # assigns an unused freshness slot number (if available) to an AEM
    # media identified by its label suffix (passed as the first argument)
    _slotfile=$(suffixtoslotfile "$1")
    rm -f "$_slotfile"

    _slotfilesglob=$(suffixtoslotfile '*')
    _lastslot=$((TPM_FRESHNESS_SLOTS - 1))
    _freeslot=$(
        {
            cat $_slotfilesglob 2>/dev/null || true
            seq 0 $_lastslot
        } | sort -n | uniq -u | head -n 1
    )

    if [ -z "$_freeslot" ]; then
        message "No more freshness token slots available!"
        return 1
    fi

    mkdir -p "${_slotfile%/*}"
    echo "$_freeslot" >> "$_slotfile"
}

checkfreshness() {
    # check whether hash of an usealed freshness token (file path
    # given as an argument) is contained in TPM NVRAM area
    _hash=$(sha1sum "$1" | cut -d ' ' -f 1)
    _lastslot=$((TPM_FRESHNESS_SLOTS - 1))
    for _i in $(seq 0 $_lastslot); do
        _slot=$(tpm_nvread_stdout -i "$TPM_FRESHNESS_INDEX" \
            -n "$((_i * 20))" -s 20 | hex)
        if [ "$_hash" == "$_slot" ]; then
            return 0
        fi
    done
    message "Freshness token does not match any slot in TPM NVRAM!"
    return 1
}

updatefreshness() {
    # takes a path to the new freshness token as an argument and
    # stores its sha1 hash in the appropriate freshness token slot
    # of the TPM NVRAM area; second argument is the AEM boot device
    # label suffix
    if [ ! -e "$TPM_FRESHNESS_PASSWORD_FILE" ]; then
        message "TPM NVRAM area AUTHWRITE password file does not exist!"
        return 1
    fi

    if ! _slot=$(suffixtoslot "$2"); then
        message "Suffix '$2' not in DB, attempting to create..."
        if ! _slot=$(assignslottosuffix "$2"); then
            message "Failed to add suffix '$2' into DB!"
            return 1
        fi
    fi

    _pw=$(cat "$TPM_FRESHNESS_PASSWORD_FILE")
    sha1sum "$1" | cut -d ' ' -f 1 | unhex \
    | tpm_nvwrite_stdin -i "$TPM_FRESHNESS_INDEX" \
      -n "$((_slot * 20))" -s 20 --password="$_pw"
}

revokefreshness() {
    # invalidates the freshness token of a specified AEM media (by its
    # label suffix
    if _slot=$(suffixtoslot "$1"); then
        message "Revoking freshness token for AEM media w/ suffix '$_suff'..."
        _pw=$(cat "$TPM_FRESHNESS_PASSWORD_FILE")
        if tpm_nvwrite -i "$TPM_FRESHNESS_INDEX" \
                -n "$((_slot * 20))" -s 20 \
                --password="$_pw" -m "0xff"; then
            message "Done."
        else
            message "Failed!"
        fi
    else
        message "AEM device with label suffix '$_suff' not found in DB!"
    fi
}

resetfreshness() {
    # invalidates ALL freshness tokens
    message "Invalidating **ALL** freshness tokens..."
    _pw=$(cat "$TPM_FRESHNESS_PASSWORD_FILE")
    if tpm_nvwrite -i "$TPM_FRESHNESS_INDEX" \
            -s "$((TPM_FRESHNESS_SLOTS * 20))" \
            --password="$_pw" -m "0xff"; then
        message "Done."
    else
        message "Failed!"
    fi
}

destroytpmnvram() {
    # releases the TPM NVRAM area; TPM owner pw as first argument
    tpm_nvrelease -i "$TPM_FRESHNESS_INDEX" --pwdo="$1"
}

synctpms() {
    _label=${1:?}
    _mnt=${2:?}

    message "Syncing to $_mnt"

    _mnt_tpms_dir=$_mnt/aem/${TPMS_DIR##*/}
    rm -rf "$_mnt_tpms_dir"

    _ids=$(ls "$TPMS_DIR")
    for _id in $_ids; do
        mkdir -p "$_mnt_tpms_dir/$_id"
        cp "$TPMS_DIR/$_id/system.data" "$_mnt_tpms_dir/$_id"

        if [ -d "$TPMS_DIR/$_id/$_label" ]; then
            cp -r  "$TPMS_DIR/$_id/$_label" "$_mnt_tpms_dir/$_id"
        fi
    done
}

devtomnt() {
    lsblk -dnr -o MOUNTPOINT "$1" 2>/dev/null |
    sed 's/%/\\x25/g' |
    xargs -0 printf
}

topdev() {
    lsblk -snrp -o KNAME "$1" | tail -n 1
}

external() {
    _aem_whole=$(topdev "$1")
    for _luks_uuid in $(getluksuuids); do
        _luks_whole=$(topdev "/dev/disk/by-uuid/$_luks_uuid")
        if [ "$_aem_whole" = "$_luks_whole" ]; then
            return 1
        fi
    done
    return 0
}

removable() {
    _rm="$(lsblk -dnr -o RM "$1") ${2-$(lsblk -dnr -o LABEL "$1")}"
    case "$_rm" in
        *.rm=[01]) _rm=${_rm##*=} ;;
                *) _rm=${_rm%% *} ;;
    esac

    [ "$_rm" = 1 ]
}
