/*
 * tee_rsa_enc_dec_TA.c
 *
 * RSA 1024 PKCS 1.5 generate key / encrypt / decrypt example
 *
 */

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

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

/* Input message */
static const char Message[] = "1111111122222222333333334444444455555555666666667777777788888888";

/* Padding label is optional */
static const char Label[] = "label";

/* RSA parameters */
#define SAMPLE_RSA_ALGO         TEE_ALG_RSAES_PKCS1_V1_5
#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_MESSAGE_SIZE  (sizeof(Message) - 1)

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

/*
 * Generates key, encrypts/decrypts message, compares plain and decrypted message
 * using RSA-1024 PKCS 1.5
 */
TEE_Result rsa_enc_dec()
{
    TEE_Result rc;
    TEE_Attribute attrs[1];
    TEE_OperationHandle hndl = NULL;
    TEE_ObjectHandle obj_hndl = NULL;
    /* Print buffer */
    char pbuf[1024] = "\0";
    uint8_t msg_encrypted[SAMPLE_RSA_SIZE];
    uint32_t msg_encrypted_len = SAMPLE_RSA_SIZE;
    uint8_t msg_decrypted[SAMPLE_MESSAGE_SIZE];
    uint32_t msg_decrypted_len = SAMPLE_MESSAGE_SIZE;

    printf("RSA encrypt/decrypt sample\n");

    printf("Generating %u-bit key\n", SAMPLE_RSA_SIZE_BIT);
    if ((rc = TEE_AllocateTransientObject(TEE_TYPE_RSA_KEYPAIR,
                                          SAMPLE_RSA_MAX_SIZE_BIT, &obj_hndl))
     || (rc = TEE_GenerateKey(obj_hndl, SAMPLE_RSA_SIZE_BIT, NULL, 0))) {
        printf("Error: key generate failed, rc = 0x%08x\n", rc);
        goto L_end;
    }

    printf("Encrypting\n");
    if ((rc = TEE_AllocateOperation(&hndl, SAMPLE_RSA_ALGO, TEE_MODE_ENCRYPT,
                                    SAMPLE_RSA_MAX_SIZE_BIT))
     || (rc = TEE_SetOperationKey(hndl, obj_hndl))) {
        printf("Error: preparation to encryption failed, rc = 0x%08x\n", rc);
        goto L_end;
    }

    TEE_InitRefAttribute(&attrs[0], TEE_ATTR_RSA_OAEP_LABEL, Label, sizeof(Label) - 1);

    if ((rc = TEE_AsymmetricEncrypt(hndl, attrs, 1, Message, SAMPLE_MESSAGE_SIZE, msg_encrypted,
                                    &msg_encrypted_len))) {
        printf("Error: encryption failed, rc = 0x%08x\n", rc);
        goto L_end;
    }
    TEE_FreeOperation(hndl);
    hndl = NULL;

    printf("Decrypting\n");
    if ((rc = TEE_AllocateOperation(&hndl, SAMPLE_RSA_ALGO, TEE_MODE_DECRYPT,
                                    SAMPLE_RSA_MAX_SIZE_BIT))
     || (rc = TEE_SetOperationKey(hndl, obj_hndl))) {
        printf("Error: preparation to decryption failed, rc = 0x%08x\n", rc);
        goto L_end;
    }

    if ((rc =  TEE_AsymmetricDecrypt(hndl, attrs, 1,msg_encrypted, msg_encrypted_len, msg_decrypted,
                                     &msg_decrypted_len))) {
        printf("Error: decryption failed, rc = 0x%08x\n", rc);
        goto L_end;
    }

    if ((msg_decrypted_len != SAMPLE_MESSAGE_SIZE)
     || (rc = TEE_MemCompare(Message, msg_decrypted, msg_decrypted_len))) {
        printf("Error: decrypted message differs, rc = 0x%08x\n", rc);
        goto L_end;
    }
    memcpy(pbuf, msg_decrypted, msg_decrypted_len);
    pbuf[msg_decrypted_len] = '\0';
    printf("Decrypted: %s\n", pbuf);

L_end:

    TEE_CloseObject(obj_hndl);
    if (hndl) {
        TEE_FreeOperation(hndl);
    }

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

    return rc;
}
