#!/bin/sh

# Exit immediately if a command exits with a non-zero status
set -e

# Summary of the script's functionality
summary="Decrypts a JWE using a TPM2.0 chip"

# TPM2.0 owner hierarchy to be used by the Operating System
auth="o"

# Display summary if requested
if [ "$1" = "--summary" ]; then
    echo "$summary"
    exit 0
fi

# Display usage information if input is from a terminal
if [ -t 0 ]; then
    exec >&2
    echo "Usage: \"zlevis-decrypt < file.jwe\""
    exit 2
fi

# Get the version of tpm2-tools
tpm2tools_version=$(tpm2_createprimary -v | awk -F'version="' '{print $2}' | awk -F'.' '{print $1}')

# Check if the tpm2-tools version is supported
if [ -z "$tpm2tools_version" ] || [ "$tpm2tools_version" -lt 4 ] || [ "$tpm2tools_version" -gt 5 ]; then
    echo "The tpm2 pin requires a tpm2-tools version between 4 and 5" >&2
    exit 1
fi

# Read the JWE protected header
read -r -d . hdr

# Decode the JWE protected header
if ! jhd="$(printf "%s" "$hdr" | jose b64 dec -i-)"; then
    echo "Error decoding JWE protected header" >&2
    exit 1
fi

# Validate the JWE pin type
if [ "$(printf "%s" "$jhd" | jose fmt -j- -Og zlevis -g pin -u-)" != "tpm2" ]; then
    echo "JWE pin mismatch" >&2
    exit 1
fi

# Extract required parameters from the JWE header
if ! hash="$(printf "%s" "$jhd" | jose fmt -j- -Og zlevis -g tpm2 -g hash -Su-)"; then
    echo "JWE missing required 'hash' header parameter" >&2
    exit 1
fi
if ! key="$(printf "%s" "$jhd" | jose fmt -j- -Og zlevis -g tpm2 -g key -Su-)"; then
    echo "JWE missing required 'key' header parameter" >&2
    exit 1
fi
if ! jwk_pub="$(printf "%s" "$jhd" | jose fmt -j- -Og zlevis -g tpm2 -g jwk_pub -Su-)"; then
    echo "JWE missing required 'jwk_pub' header parameter" >&2
    exit 1
fi
if ! jwk_priv="$(printf "%s" "$jhd" | jose fmt -j- -Og zlevis -g tpm2 -g jwk_priv -Su-)"; then
    echo "JWE missing required 'jwk_priv' header parameter" >&2
    exit 1
fi

# Handle optional PCR parameters
pcr_ids="$(printf "%s" "$jhd" | jose fmt -j- -Og zlevis -g tpm2 -g pcr_ids -Su-)" || true
pcr_spec=""
if [ -n "$pcr_ids" ]; then
    pcr_bank="$(printf "%s" "$jhd" | jose fmt -j- -Og zlevis -g tpm2 -g pcr_bank -Su-)"
    pcr_spec="$pcr_bank:$pcr_ids"
fi

# Define and trap tmp jwk_pub and jwk_priv
tmp_jwk_pub="/tmp/jwk_pub.$$"
tmp_jwk_priv="/tmp/jwk_priv.$$"
trap 'rm -f "$tmp_jwk_pub" "$tmp_jwk_priv"' EXIT

# Decode the public and private keys from Base64
if ! printf "%s" "$jwk_pub" | jose b64 dec -i- -O "$tmp_jwk_pub"; then
    echo "Decoding jwk.pub from Base64 failed" >&2
    exit 1
fi
if ! printf "%s" "$jwk_priv" | jose b64 dec -i- -O "$tmp_jwk_priv"; then
    echo "Decoding jwk.priv from Base64 failed" >&2
    exit 1
fi

# Define and trap primary_context
tmp_primary_context="/tmp/primary_context.$$"
trap 'rm -f "$tmp_jwk_pub" "$tmp_jwk_priv" "$tmp_primary_context"' EXIT

# Create the primary key in the TPM
case "$tpm2tools_version" in
    4|5) tpm2_createprimary -Q -C "$auth" -g "$hash" -G "$key" -c "$tmp_primary_context" || fail=$?;;
    *) fail=1;;
esac
if [ -n "$fail" ]; then
    echo "Creating TPM2 primary key failed" >&2
    exit 1
fi
tpm2_flushcontext -t

# Define and trap load_context
tmp_load_context="/tmp/load_context.$$"
trap 'rm -f "$tmp_jwk_pub" "$tmp_jwk_priv" "$tmp_primary_context" "$tmp_load_context"' EXIT

# Load the JWK into the TPM
case "$tpm2tools_version" in
    4|5) tpm2_load -Q -C "$tmp_primary_context" -u "$tmp_jwk_pub" -r "$tmp_jwk_priv" -c "$tmp_load_context" || fail=$?;;
    *) fail=1;;
esac
if [ -n "$fail" ]; then
    echo "Loading jwk to TPM2 failed" >&2
    exit 1
fi
tpm2_flushcontext -t

# Remove tmp_jwk_pub, tmp_jwk_priv and tmp_primary_context
rm -f "$tmp_jwk_pub" "$tmp_jwk_priv" "$tmp_primary_context"

# Unseal the JWK from the TPM
case "$tpm2tools_version" in
    4|5) jwk="$(tpm2_unseal -c "$tmp_load_context" ${pcr_spec:+-p pcr:$pcr_spec})" || fail=$?;;
    *) fail=1;;
esac
if [ -n "$fail" ]; then
    echo "Unsealing jwk from TPM failed" >&2
    exit 1
fi
tpm2_flushcontext -t

# Remove tmp_load_context
rm -f "$tmp_load_context"

# Output the decrypted JWK along with the original header
(echo "$jwk$hdr."; /bin/cat) | jose jwe dec -k- -i-

# Exit with the status of the last command
exit $?