#!/bin/bash
set -e
usage() {
    echo "" 1>&2;
    echo "This script is used to sign a file with a TA Authority private key" 1>&2;
    echo "Usage: $0 <options> <file> <output_file>" 1>&2;
    echo "" 1>&2;
    echo " <file> - Location of the file to sign." 1>&2;
    echo " <output_file> - Output file name and location (default <file>.signed)" 1>&2;
    echo " <options>" 1>&2;
    echo "  -t <pkg type>   Package type: SEC1(deprecated), SEC2(default), SEC3(version anti-rollback)," 1>&2;
    echo "                                SEC4(version anti-rollback + encryption)" 1>&2;
    echo "  -c <location>   Override location of certificate (default ./ta_auth/cert.pem)" 1>&2;
    echo "  -k <location>   Override location of private key (default ./ta_auth/private/key.pem)" 1>&2;
    echo "  -p <password>   Password to unlock the private key" 1>&2;
    echo "  -e <name>       Use OpenSSL engine" 1>&2;
    echo "  -r <version>    TA anti-rollback version (non-zero, valid for SEC3 and SEC4)" 1>&2;
    echo "  -1 <ta key>     AES key to encrypt TA (SEC4 only)" 1>&2;
    echo "  -2 <auth key>   AES key to encrypt TA key (SEC4 only)" 1>&2;
    echo "  -n <name>       Chipset name (needed for use HSM)" 1>&2;
    echo "  -a              Append certificate to the signed file (optional, only for backward compatibility)" 1>&2;
    echo "  -h              Display this usage information" 1>&2;
    exit 1;
}

CERT=""
KEY=""
PKG_TYPE="SEC2"
APPEND=""
PASS=""
ENGINE=""
VERSION=""
TA_KEY=""
AUTH_KEY=""
CHIPSET_NAME=""
HSM_ENABLED="false"
while getopts "t:c:k:p:e:n:r:1:2:ah" o; do
    case "${o}" in
        c)
            CERT=${OPTARG}
            ;;
        k)
            KEY=${OPTARG}
            ;;
        p)
            PASS="-passin pass:${OPTARG}"
            ;;
        e)
            ENGINE="-engine ${OPTARG} -keyform ENGINE"
            ;;
        t)
            PKG_TYPE=${OPTARG}
            ;;
        r)
            VERSION=${OPTARG}     
            ;;
        1)
            TA_KEY=${OPTARG}
            ;;
        2)
            AUTH_KEY=${OPTARG}
            ;;
        a)
            APPEND="true"
            ;;
        n)
            CHIPSET_NAME="${OPTARG}"
            ;;
        h)
            usage
            ;;
    esac
done
shift $((OPTIND-1))

if [[ -z "$1" ]]; then
    usage
fi

if [[ -z "$3" ]]; then
    usage
fi
PUB_KEY=$3

if [[ ! -f "$1" ]]; then
    echo "File $1 does not exist"
    exit 1;
fi

if [[ ! -f "$CERT" ]]; then
    echo "Certificate $CERT does not exist"
    exit 1;
fi

if [[ ! -f "$KEY" ]]; then
    echo "Private key $KEY does not exist"
    exit 1;
fi

if [ -n "$ENGINE" ]; then #HSM enanbled
    if [ "$CHIPSET_NAME" == "" ]; then
        echo "For use HSM, Chipset name (option -n) is required"
        usage
    fi

    HSM_ENABLED="true"
    echo "HSM module enabled"
fi

case "$PKG_TYPE" in
    "SEC1")
        echo "SEC1 is deprecated"
        exit 1;
        ;;
    "SEC2")
        if [ -n "$VERSION" ]; then
            echo "Ignoring TA version parameter"
        fi
        if [ -n "$AES_KEY" ]; then
            echo "Ignoring AES key parameter"
        fi
        if [ -n "$RSA_KEY" ]; then
            echo "Ignoring RSA key parameter"
        fi
        ;;
    "SEC3")
        if [ -z "$VERSION" ] || [[ ! $VERSION -gt 0 ]]; then
            echo "SEC3 requires version parameter 1 or more (-r)!"
            exit 1;
        fi
        if [ -n "$TA_KEY" ]; then
            echo "Ignoring TA key parameter"
        fi
        if [ -n "$AUTH_KEY" ]; then
            echo "Ignoring auth key parameter"
        fi
        ;;
    "SEC4")
        if [ -z "$VERSION" ] || [[ ! $VERSION -gt 0 ]]; then
            echo "SEC4 requires version parameter 1 or more (-r)!"
            exit 1;
        fi
        if [ ! -f "$TA_KEY" ]; then
            echo "SEC4 requires TA key parameter (-1)"
            exit 1;
        fi
        if [ ! -f "$AUTH_KEY" ]; then
            echo "SEC4 requires auth key parameter (-2)"
            exit 1;
        fi
        ;;
esac

OUTPUT="$1.signed"
UNSIGNED="$1.unsigned.$$.$RANDOM"
if [[ ! -z "$2" ]]; then
    OUTPUT=$2
fi

# size of the original ELF. For SEC4, it's going to be overrided by the encrypted image len
SIZE=$(stat -c%s "$1")

if [ "$PKG_TYPE" == "SEC2" ]; then
    # for SEC2, only ELF is signed
    cat $1 > $UNSIGNED
elif [ "$PKG_TYPE" == "SEC3" ]; then
    # for SEC3, elf and version are signed
    cat $1 > $UNSIGNED
    # version len and value
    echo "Version: $VERSION"
    printf "0: %.8x" 4 | /usr/bin/xxd -r -g0 >> $UNSIGNED
    printf "0: %.8x" $VERSION | /usr/bin/xxd -r -g0 >> $UNSIGNED
elif [ "$PKG_TYPE" == "SEC4" ]; then
    TA_IV=$(/usr/bin/xxd -l 16 -ps /dev/urandom)
    TA_KEY_BYTES=$(/usr/bin/xxd -c32 -ps $TA_KEY)
    AUTH_IV=$(/usr/bin/xxd -l 16 -ps /dev/urandom)
    AUTH_KEY_BYTES=$(/usr/bin/xxd -c32 -ps $AUTH_KEY)
    if [[ ($(stat -c%s "$TA_KEY") -ne 32) || ($(stat -c%s "$AUTH_KEY") -ne 32) ]]; then
        echo "Invalid size of TA_KEY or AUTH_KEY"
        exit 1
    fi
    echo "TA_IV: $TA_IV"
    echo "TA_KEY: $TA_KEY_BYTES"
    echo "AUTH_IV: $AUTH_IV"
    echo "AUTH_KEY: $AUTH_KEY_BYTES"
    
    # encrypted ELF
    openssl enc -aes-256-ctr -K $TA_KEY_BYTES -iv $TA_IV -in $1 > $UNSIGNED
    
    # calculate encrypted ELF size
    SIZE=$(stat -c%s "$UNSIGNED")

    # version len and value
    echo "Version: $VERSION"
    printf "0: %.8x" 4 | /usr/bin/xxd -r -g0 >> $UNSIGNED
    printf "0: %.8x" $VERSION | /usr/bin/xxd -r -g0 >> $UNSIGNED
    
    # TA IV len and value
    printf "0: %.8x" 16 | /usr/bin/xxd -r -g0 >> $UNSIGNED
    printf "0: $TA_IV" | /usr/bin/xxd -r -g0 >> $UNSIGNED
    
    # AUTH IV len and value
    printf "0: %.8x" 16 | /usr/bin/xxd -r -g0 >> $UNSIGNED
    printf "0: $AUTH_IV" | /usr/bin/xxd -r -g0 >> $UNSIGNED

    # sha-1 checksum of the cleartext ELF
    printf "0: %.8x" 20 | /usr/bin/xxd -r -g0 >> $UNSIGNED
    openssl sha1 -binary $1 >> $UNSIGNED
    echo "SHA-1: $(/usr/bin/xxd -ps -s-20 $UNSIGNED)"
    
    # encrypted AES key len and value
    printf "0: %.8x" 32 | /usr/bin/xxd -r -g0 >> $UNSIGNED    
    openssl enc -aes-256-ctr -K $AUTH_KEY_BYTES -iv $AUTH_IV -in $TA_KEY >> $UNSIGNED
else
    echo "Unknown pkg type $PKG_TYPE"
    exit 1    
fi

# starting to write package file...

# magic
echo -n "$PKG_TYPE" > $OUTPUT

# ELF size
printf "0: %.8x" $SIZE | /usr/bin/xxd -r -g0 >> $OUTPUT

# ELF (+version for SEC3, +version +key for SEC4)
cat $UNSIGNED >> $OUTPUT

# CN len and value
# subject= /DC=debug/CN=dev_ta:dev_ta
# TODO: [a-z,A-Z,0-9] should follow regular expression matching to SM-* (model name)
CN=`openssl x509 -in $CERT -noout -subject | tr -d "[:blank:]" | sed 's#subject=.*CN=##g'`
if [ -z $CN ]; then
    echo "fail to get CN = $CN"
    exit 1
else
    echo "get CN = $CN"
fi
printf "0: %.2x" ${#CN}  | /usr/bin/xxd -r -g0 >> $OUTPUT
echo -n $CN >> $OUTPUT

# signature
if [ -n "$KEY" ]; then
    if openssl rsa -in $KEY $PASS -noout &>/dev/null; then
        KEY_ALGO="rsa"
    elif openssl ec -in $KEY $PASS -noout &>/dev/null; then
        KEY_ALGO="ec"
    else
        echo "Unknown key type"
        exit 1
    fi
fi

if [ "$HSM_ENABLED" == "true" ]; then #HSM enanbled
    source hsm_setup.sh

    #Environment configs of HSM
    CHIPSET_NAME=`echo ${CHIPSET_NAME,,} | sed 's/[^a-z,A-Z,0-9]//g'`  #HSM only support lowcase key index name
    HSM_KEY_INDEX="${HSM_KEY_INDEX_PREFIX}${CHIPSET_NAME}${KEY_TYPE}${HSM_KEY_INDEX_SUFFIX}"
    echo "HSM Key index: ${HSM_KEY_INDEX}"

    OPENSSL_CMD=${HSM_NCIPHER_BIN_PATH}${HSM_OPENSSL_CMD}
    KEY=${HSM_KEY_INDEX}
else
    OPENSSL_CMD=openssl
    # with RSA certificate, SEC3 and higher require PSS
    if [ "$KEY_ALGO" == "rsa" ] && [ "$PKG_TYPE" != "SEC2" ]; then
        SIGN_PARAM="-sigopt rsa_padding_mode:pss -sigopt rsa_pss_saltlen:-1"
    fi
fi

SIGNATURE=$(mktemp)
$OPENSSL_CMD dgst -sha256 $SIGN_PARAM $PASS $ENGINE -sign $KEY -binary $UNSIGNED > ${SIGNATURE}
BYTES=`cat ${SIGNATURE} | /usr/bin/wc --bytes`
printf "0: %.4x" $BYTES | /usr/bin/xxd -r -g0 >> $OUTPUT
	
if [[ ! -f "$PUB_KEY" ]]; then
    $OPENSSL_CMD x509 -in $CERT -pubkey -noout > $PUB_KEY
fi

if ! $OPENSSL_CMD dgst -sha256 $SIGN_PARAM -verify $PUB_KEY -signature $SIGNATURE $UNSIGNED; then
    echo "Fail to verify signature"
    exit 1
fi

cat ${SIGNATURE} >> $OUTPUT
rm ${SIGNATURE}

# certificate
BYTES=`openssl x509 -in $CERT -inform PEM -outform DER | /usr/bin/wc --bytes`
if [ $BYTES -le 0 ]; then
    echo "fail to get BYTES : $BYTES"
    exit 1
else
    echo "get BYTES : $BYTES"
fi
printf "0: %.4x" $BYTES | /usr/bin/xxd -r -g0 >> $OUTPUT
if openssl x509 -in $CERT -inform PEM -outform DER >> $OUTPUT; then
    echo "add certificate"
else
    echo "fail to add certificate"
    exit 1
fi

# cleanup
rm $UNSIGNED