/*
 * tee_rsa_sign_verif_TA.c
 *
 * RSA-1024 PSS SHA256 sign/verify example
 *
 * NOTE: As test data is initially presented in form of hex-string the additional
 *       hex-string to byte-array conversion is used.
 */

#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <tee_internal_api.h>
#include "crypto_samples_functions.h"

/*************************************************************************************************/

/* Modulus */
static const char N[] = "a8d68acd413c5e195d5ef04e1b4faaf242365cb450196755e92e1215ba59802a"
                        "afbadbf2564dd550956abb54f8b1c917844e5f36195d1088c600e07cada5c080"
                        "ede679f50b3de32cf4026e514542495c54b1903768791aae9e36f082cd38e941"
                        "ada89baecada61ab0dd37ad536bcb0a0946271594836e92ab5517301d45176b5";
/* Public exponent*/
static const char E[] = "11";
/* Private exponent*/
static const char D[] = "04f7405154be02c482bf07115b27fd7f988920d8207937bec317f1791c11b0f2"
                        "325851c35ce42bebc828c946438cc22656b702cce2a0dad6d8a5ac3fe6fd587c"
                        "36f81ff5edc977edf9c6c085efd73510ddb2532742367f8f878814c1714549dc"
                        "fa17dbd5d1e17ce3dd1a9893a3785bc34c5205df3978071aa36c9e86bd33cba1";
/* Primes 1, 2*/
static const char P[] = "c107a2fe924b76e206cb9bc4af2ab7008547c00846bf6d0680b3eac3ebcbd0c7"
                        "fd7a54c2b9899b08f80cde1d3691eaaa2816b1eb11822d6be7beaf4e30977c49";
static const char Q[] = "dfea984ce4307eafc0d140c2bb82861e5dbac4f8567cbc981d70440dd6394920"
                        "79031486315e305eb83e591c4a2e96064966f7c894c3ca351925b5ce82d8ef0d";
/* Exponents 1, 2*/
static const char DP[] = "2210680ece85c9af6a9c66c85b25a7d2ea57f4b62a9a406a8f2ecf1383f6ca7d"
                         "a533b49ad572a2e37711545f8219bfffe8f4f2388a9e807c7430b58644cf7049";
static const char DQ[] = "b866d7c6da09d1bdea15bcdc9a6b7d826b4e8417cec113c890b6cea1fbb6b4b1"
                         "54991ff60a89cd7b2e51768fc49ed5e70f45bcff89922e0d9c3d2c4fb70d0119";
/* Coefficient */
static const char QP[] = "81aa7cc8ebea78b312d58e5b8047c2ce77e067af07eed1f86a7fec2479956435"
                         "2f3864ee209307e5c0673a0d557cb754807072ab752cacc2f874f5938890a8da";

/* Message digest */
static const char Digest[] = "7936c953ec5781739403588d1e04380742e054a1b5eede9342e854e579975ca2";

/* Pre-computed signature */
static const char Signature[] = "938b28be99e59e79d1004667095d9ebe0f1563693eea80d368a793fad5d766d5"
                                "019d0b78e27b9007d32b3a4d072f75bcce49e5e2c442919785c5c0b2242434de"
                                "1b4bcaf6b0cce74c53881718f7a23be9c939acbcc2a3793df0a8f2ee7c73d743"
                                "f74a29262a211aef0841a73929678fe390edef428c2a3d2c65ed19d226ab2ea4";
/* SHA256 message digest size */
#define SAMPLE_SHA_DIGEST_SIZE 32

/* RSA parameters */
#define SAMPLE_RSA_ALGO         TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256
#define SAMPLE_RSA_SIZE_BIT     1024
#define SAMPLE_RSA_SIZE         (SAMPLE_RSA_SIZE_BIT / 8)
#define SAMPLE_RSA_MAX_SIZE_BIT SAMPLE_RSA_SIZE_BIT
#define SAMPLE_RSA_N_SIZE       ((sizeof(N) - 1) / 2)
#define SAMPLE_RSA_E_SIZE       ((sizeof(E) - 1) / 2)
#define SAMPLE_RSA_D_SIZE       ((sizeof(D) - 1) / 2)
#define SAMPLE_RSA_P_SIZE       ((sizeof(P) - 1) / 2)
#define SAMPLE_RSA_Q_SIZE       ((sizeof(Q) - 1) / 2)
#define SAMPLE_RSA_DP_SIZE      ((sizeof(DP) - 1) / 2)
#define SAMPLE_RSA_DQ_SIZE      ((sizeof(DQ) - 1) / 2)
#define SAMPLE_RSA_QP_SIZE      ((sizeof(QP) - 1) / 2)

/*************************************************************************************************/

/*
 * Signs and verifies message digest using RSA-1024 PSS
 */
TEE_Result rsa_sign_verify(void)
{
    TEE_Result rc;
    TEE_Attribute attrs[8];
    TEE_OperationHandle hndl_sign = NULL;
    TEE_OperationHandle hndl_verify = NULL;
    TEE_ObjectHandle obj_hndl_sign = NULL;
    TEE_ObjectHandle obj_hndl_verify = NULL;
    /* Print buffer */
    char pbuf[1024] = "\0";
    /* Message digest: Digest */
    uint8_t msg_digest[SAMPLE_SHA_DIGEST_SIZE];
    /* RSA key parameters */
    uint8_t n[SAMPLE_RSA_N_SIZE];
    uint8_t e[SAMPLE_RSA_E_SIZE];
    uint8_t d[SAMPLE_RSA_D_SIZE];
    uint8_t p[SAMPLE_RSA_P_SIZE];
    uint8_t q[SAMPLE_RSA_Q_SIZE];
    uint8_t dp[SAMPLE_RSA_DP_SIZE];
    uint8_t dq[SAMPLE_RSA_DQ_SIZE];
    uint8_t qp[SAMPLE_RSA_QP_SIZE];
    /* Pre-computed message signature: Signature */
    uint8_t msg_signature_precomputed[SAMPLE_RSA_SIZE];
    /* Message signature */
    uint8_t msg_signature[SAMPLE_RSA_SIZE];
    uint32_t msg_signature_len = SAMPLE_RSA_SIZE;

    printf("RSA sign/verify sample\n");

    /* Prepare data */
    if (hex_string2bytes(Digest, msg_digest, SAMPLE_SHA_DIGEST_SIZE)
     || hex_string2bytes(N, n, SAMPLE_RSA_N_SIZE)
     || hex_string2bytes(E, e, SAMPLE_RSA_E_SIZE)
     || hex_string2bytes(D, d, SAMPLE_RSA_D_SIZE)
     || hex_string2bytes(P, p, SAMPLE_RSA_P_SIZE)
     || hex_string2bytes(Q, q, SAMPLE_RSA_Q_SIZE)
     || hex_string2bytes(DP, dp, SAMPLE_RSA_DP_SIZE)
     || hex_string2bytes(DQ, dq, SAMPLE_RSA_DQ_SIZE)
     || hex_string2bytes(QP, qp, SAMPLE_RSA_QP_SIZE)) {
        printf("Error: data preparation failed\n");
        return TEE_ERROR_GENERIC;
    }

    printf("Signing\n");

    TEE_InitRefAttribute(&attrs[0], TEE_ATTR_RSA_MODULUS, n, SAMPLE_RSA_N_SIZE);
    TEE_InitRefAttribute(&attrs[1], TEE_ATTR_RSA_PUBLIC_EXPONENT, e, SAMPLE_RSA_E_SIZE);
    TEE_InitRefAttribute(&attrs[2], TEE_ATTR_RSA_PRIVATE_EXPONENT, d, SAMPLE_RSA_D_SIZE);
    TEE_InitRefAttribute(&attrs[3], TEE_ATTR_RSA_PRIME1, p, SAMPLE_RSA_P_SIZE);
    TEE_InitRefAttribute(&attrs[4], TEE_ATTR_RSA_PRIME2, q, SAMPLE_RSA_Q_SIZE);
    TEE_InitRefAttribute(&attrs[5], TEE_ATTR_RSA_EXPONENT1, dp, SAMPLE_RSA_DP_SIZE);
    TEE_InitRefAttribute(&attrs[6], TEE_ATTR_RSA_EXPONENT2, dq, SAMPLE_RSA_DQ_SIZE);
    TEE_InitRefAttribute(&attrs[7], TEE_ATTR_RSA_COEFFICIENT, qp, SAMPLE_RSA_QP_SIZE);

    if ((rc = TEE_AllocateTransientObject(TEE_TYPE_RSA_KEYPAIR, SAMPLE_RSA_MAX_SIZE_BIT,
                                                         &obj_hndl_sign))
     || (rc = TEE_AllocateOperation(&hndl_sign, SAMPLE_RSA_ALGO, TEE_MODE_SIGN,
                                    SAMPLE_RSA_MAX_SIZE_BIT))
     || (rc = TEE_PopulateTransientObject(obj_hndl_sign, attrs, 8))
     || (rc = TEE_SetOperationKey(hndl_sign, obj_hndl_sign))) {
        printf("Error: preparation to signing failed, rc = 0x%08x\n", rc);
        goto L_end;
    }

    if ((rc = TEE_AsymmetricSignDigest(hndl_sign, NULL, 0, msg_digest, SAMPLE_SHA_DIGEST_SIZE,
                                       msg_signature, &msg_signature_len))) {
        printf("Error: digest signing failed, rc = 0x%08x\n", rc);
        goto L_end;
    }

    bytes2hex_string(pbuf, msg_signature, msg_signature_len);
    printf("Signature: %s\n", pbuf);

    printf("Verifying results of signing operation\n");

    TEE_InitRefAttribute(&attrs[0], TEE_ATTR_RSA_MODULUS, n, SAMPLE_RSA_N_SIZE);
    TEE_InitRefAttribute(&attrs[1], TEE_ATTR_RSA_PUBLIC_EXPONENT, e, SAMPLE_RSA_E_SIZE);

    if ((rc = TEE_AllocateTransientObject(TEE_TYPE_RSA_PUBLIC_KEY, SAMPLE_RSA_MAX_SIZE_BIT,
                                          &obj_hndl_verify))
     || (rc = TEE_AllocateOperation(&hndl_verify, SAMPLE_RSA_ALGO, TEE_MODE_VERIFY,
                                    SAMPLE_RSA_MAX_SIZE_BIT))
     || (rc = TEE_PopulateTransientObject(obj_hndl_verify, attrs, 2))
     || (rc = TEE_SetOperationKey(hndl_verify, obj_hndl_verify))) {
        printf("Error: preparation to verification failed, rc = 0x%08x\n", rc);
        goto L_end;
    }

    if ((rc = TEE_AsymmetricVerifyDigest(hndl_verify, NULL, 0, msg_digest, SAMPLE_SHA_DIGEST_SIZE,
                                         msg_signature, SAMPLE_RSA_SIZE))) {
        printf("Error: digest verification failed, rc = 0x%08x\n", rc);
        goto L_end;
    }
    printf("Verified. OK.\n");

    printf("Verifying pre-computed signature\n");

    if (hex_string2bytes(Signature, msg_signature_precomputed, SAMPLE_RSA_SIZE)) {
        printf("Error: wrong format of pre-computed signature\n");
        rc = TEE_ERROR_GENERIC;
        goto L_end;
    }

    bytes2hex_string(pbuf, msg_signature_precomputed, SAMPLE_RSA_SIZE);
    printf("Signature: %s\n", pbuf);

    if ((rc = TEE_AsymmetricVerifyDigest(hndl_verify, NULL, 0, msg_digest, SAMPLE_SHA_DIGEST_SIZE,
                                         msg_signature_precomputed, SAMPLE_RSA_SIZE))) {
        printf("Error: pre-computed digest verification failed, rc = 0x%08x\n", rc);
        goto L_end;
    }
    printf("Verified. OK.\n");

L_end:

    TEE_CloseObject(obj_hndl_sign);
    if (hndl_sign) {
        TEE_FreeOperation(hndl_sign);
    }
    TEE_CloseObject(obj_hndl_verify);
    if (hndl_verify) {
        TEE_FreeOperation(hndl_verify);
    }

    if (rc) {
        printf("RSA sign/verify failed\n");
    } else {
        printf("RSA sign/verify passed\n");
    }

    return rc;
}
