#!/bin/sh

set -e

FORCE=0
READ_KEY=""
WRITE_KEY=""
OUTPUT_BINARY=0
ROW_COUNT=8
ROW_OFFSET=0

die() {
    echo "$@" >&2
    exit 1
}

usage() {
    cat <<EOF
   $(basename "$0") [-cfwy] <key>

   No args - reads the current device unique private key from OTP.
   *These values are NOT visible via 'vcgencmd otp_dump'*

   -b Output the key in binary format.
   -c Reads key and exits with 1 if it is all zeros i.e. not set.
   -f Force write (if OTP is non-zero).
      The vcmailbox API checks that the new key is equal to the bitwise OR of the current OTP and the new key.
      N.B. OTP bits can never change from 1 to 0.
   -w Writes the new key to OTP memory.
   -y Skip the confirmation prompt when writing to OTP.
   -l Specify key length in words. Defaults to 8 words (32 bytes). Pi 5 supports up to 16 words (64 bytes).
   -o word Offset into the keystore to use, e.g. 0-7 for Pi 4, 0-15 for Pi 5. Defaults to zero.

   <key> is usually a 64 digit hex number (256 bit)

   Key generation:
   The Raspberry Pi firmware cryptography services requires that the device unique private key is
   a valid ECDSA with P-256 curve key. Due to limited OTP space only the raw private key component (d)
   is stored in OTP.

   Example key generation and provisioning:

   # Generate the new private-key - remember to save this to a secure, off-device location!
   openssl ecparam -name prime256v1 -genkey -noout -out private_key.pem

   # Extract raw the private key component
   openssl ec -in private_key.pem -text -noout | awk '/priv:/{flag=1; next} /pub:/{flag=0} flag' | tr -d ' \n:' | head -n1 > d.hex

   # Write the key to OTP
   rpi-otp-private-key -w \$(cat d.hex)

   IMPORTANT: Raspberry Pi 5 and earlier revisions do not have a hardware secure key store
   so the raw OTP values are potentially readable by processes with root-privileges.

   In newer firmware releases, the mailbox APIs used by this script to read the OTP can
   be disabled by setting lock_device_private_key=1 in config.txt.
   On Pi4 or newer, if secure-boot is enabled, then this parameter cannot be
   tampered with because config.txt is stored within the signed boot.img.

   WARNING: Changes to OTP memory are permanent and cannot be undone.
EOF
exit 1
}

check_key_set() {
    read_key
    if [ -z "$(echo "${READ_KEY}" | sed s/0//g)" ]; then
        return 1
    fi
    return 0
}

read_key() {
    out=READ_KEY="$(vcmailbox 0x00030081 $((8 + ROW_COUNT * 4)) $((8 + ROW_COUNT * 4)) $ROW_OFFSET $ROW_COUNT 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0)" || die "Failed to read the current key from OTP"
    READ_KEY="$(echo "${out}" | sed 's/0x//g' | awk -v last=$((8 + ROW_COUNT)) '{for(i=8;i<last;i++) printf $i; print ""}')"
}

write_key() {
    key="${1}"
    # Normalize formatting and check the length
    key="$(echo "${key}" | tr 'A-Z' 'a-z')"
    key="$(echo "${key}" | sed 's/[^a-f0-9]//g')"
    [ "$(echo -n "${key}" | wc -c)" = $((ROW_COUNT * 4 * 2)) ] || die "Invalid key parameter"

    count=0
    key_params=""
    while [ ${count} -lt $ROW_COUNT ]; do
        start=$(((count * 8) + 1))
        end=$((start + 7))
        key_params="${key_params} 0x$(echo -n "${key}" | cut -c${start}-${end})"
        count=$((count + 1))
    done

    if [ "${YES}" = 0 ] && [ -t 0 ]; then
       echo "Write ${key} of $ROW_COUNT words to OTP starting at word offset $ROW_OFFSET?"
       echo
       echo "WARNING: Updates to OTP registers are permanent and cannot be undone."

       echo "Type YES (in upper case) to continue or press return to exit."
       read -r confirm
       if [ "${confirm}" != "YES" ]; then
          echo "Cancelled"
          exit
       fi
    fi

    vcmailbox 0x38081 $((8 + ROW_COUNT * 4)) $((8 + ROW_COUNT * 4)) $ROW_OFFSET $ROW_COUNT ${key_params} || die "Failed to write key"
    read_key
    [ "${READ_KEY}" = "${key}" ] || die "Key readback check failed. ${out}"
}

YES=0
while getopts bcfhw:yl:o: option; do
    case "${option}" in
        b) OUTPUT_BINARY=1
            ;;
        c)
            if check_key_set; then
                exit 0
            fi
            exit 1
            ;;
        f) FORCE=1
            ;;
        h) usage
            ;;
        w) WRITE_KEY="${OPTARG}"
            ;;
        y) YES=1
            ;;
        l) ROW_COUNT="${OPTARG}"
            ;;
        o) ROW_OFFSET="${OPTARG}"
            ;;
        *) echo "Unknown argument \"${option}\""
            usage
            ;;
    esac
done

if [ -f "/sys/firmware/devicetree/base/system/linux,revision" ]; then
    BOARD_INFO="$(od -v -An -t x1 /sys/firmware/devicetree/base/system/linux,revision | tr -d ' \n')"
elif grep -q Revision /proc/cpuinfo; then
    BOARD_INFO="$(sed -n '/^Revision/s/^.*: \(.*\)/\1/p' < /proc/cpuinfo)"
elif command -v vcgencmd > /dev/null; then
    BOARD_INFO="$(vcgencmd otp_dump | grep '30:' | sed 's/.*://')"
else
    die "No Raspberry Pi board info found"
fi
if [ $(((0x$BOARD_INFO >> 23) & 1)) = 0 ]; then
    die "Chip not supported"
fi
if [ $(((0x$BOARD_INFO >> 12) & 15)) = 3 ]; then
    MAX_ROW_COUNT=8
elif [ $(((0x$BOARD_INFO >> 12) & 15)) = 4 ]; then
    MAX_ROW_COUNT=16
else
    echo "WARNING: Secure-boot is only supported on Pi4 and newer models"
    MAX_ROW_COUNT=8
fi
if [ -z "$ROW_COUNT" ] || [ "$ROW_COUNT" -ne "$ROW_COUNT" ] 2>/dev/null; then
    die "Key length not a number"
fi
if [ $ROW_COUNT -lt 1 ]; then
    die "Length too small"
fi
if [ $ROW_COUNT -gt $MAX_ROW_COUNT ]; then
    die "Length too big"
fi
if [ -z "$ROW_OFFSET" ] || [ "$ROW_OFFSET" -ne "$ROW_OFFSET" ] 2>/dev/null; then
    die "Offset is not a number"
fi
if [ $ROW_OFFSET -lt 0 ]; then
    die "Offset too small"
fi
if [ ${ROW_OFFSET} -gt $((MAX_ROW_COUNT - ROW_COUNT)) ]; then
    die "Offset too big"
fi
if [ -z $(which vcmailbox) ]; then
    die "vcmailbox command missing"
fi
if [ -z $(which xxd) ] && [ "$OUTPUT_BINARY" -eq "1" ]; then
    die "xxd command missing"
fi

if [ -n "${WRITE_KEY}" ]; then
    if [ "${FORCE}" = 0 ] && check_key_set; then
        die "Current key is non-zero. Specify -f to write anyway"
    fi
    write_key "${WRITE_KEY}"
else
    read_key
    if [ "${OUTPUT_BINARY}" = 1 ]; then
        echo "${READ_KEY}" | xxd -r -p
    else
        echo "${READ_KEY}"
    fi
fi