/*
 * tee_aes_TA.c
 *
 * AES-128 GCM SIV encrypt/decrypt example
 *
 * NOTE: Test data was obtained from BoringSSL project:
 *       boringssl/src/crypto/fipsmodule/modes/gcm_test.cc.
 *       As this 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"
#include <openssl/base.h>
#include <openssl/err.h>

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

/* Secret key */
static const char KEY[] = "feffe9928665731c6d6a8f9467308308";
/* IV */
static const char NONCE[] = "cafebabefacedbaddecaf888";
/* Plain text */
static const char IN[] = "d9313225f88406e5a55909c5aff5269a86a7a9531534f7da2e4c303d8a31"
                         "8a721c3c0c95956809532fcf0e2449a6b525b16aedf5aa0de657ba637b39";
/* A data */
static const char AD[] = "feedfacedeadbeeffeedfacedeadbeefabaddad2";
/* Cipher text */
static const char CT[] = "42831ec2217774244b7221b784d0d49ce3aa212f2c02a4e035c17e2329ac"
                         "a12e21d514b25466931c7d8f6a5aac84aa051ba30b396a0aac973d58e091";
/* Tag */
static const char TAG[] = "5bc94fbc3221a5db94fae95ae7121a47";

/* AES parameters */
#define SAMPLE_AES_ALGO         TEE_ALG_AES_GCM

#define SAMPLE_AES_SIZE_BIT     128
#define SAMPLE_AES_SIZE         (SAMPLE_AES_SIZE_BIT / 8)

#define SAMPLE_AES_MAX_SIZE_BIT SAMPLE_AES_SIZE_BIT

#define SAMPLE_AES_KEY_SIZE     ((sizeof(KEY) - 1) / 2)

#define SAMPLE_AES_IV_SIZE      ((sizeof(NONCE) - 1) / 2)
#define SAMPLE_AES_AD_SIZE      ((sizeof(AD) - 1) / 2)
#define SAMPLE_AES_TAG_SIZE     ((sizeof(TAG) - 1) / 2)

/* Message parameters */
#define SAMPLE_MSG_SIZE         ((sizeof(IN) - 1) / 2)
#define SAMPLE_ENC_MSG_SIZE     ((SAMPLE_MSG_SIZE > SAMPLE_AES_SIZE) ?\
                                  SAMPLE_MSG_SIZE : SAMPLE_AES_SIZE)
#define SAMPLE_CT_SIZE          ((sizeof(CT) - 1) / 2)

/* Defines message path length for encrypt update example */
#define SAMPLE_MSG_INITIAL_LEN 5

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

/*
 * Encode / decode message using AES-128 GCM SIV and compares results with pre-calculated data
 */
TEE_Result aes_enc_dec(void)
{
    TEE_Result rc;
    TEE_Attribute attrs[1];
    TEE_OperationHandle hndl = NULL;
    TEE_ObjectHandle obj_hndl = NULL;
    /* Print buffer */
    char pbuf[1024] = "\0";
    /* AES parameters */
    uint8_t key[SAMPLE_AES_KEY_SIZE];
    uint8_t iv[SAMPLE_AES_IV_SIZE];
    uint8_t ad[SAMPLE_AES_AD_SIZE];
    uint8_t tag[SAMPLE_AES_TAG_SIZE];
    uint32_t tag_len = SAMPLE_AES_TAG_SIZE;
    uint8_t tag_precomputed[SAMPLE_AES_TAG_SIZE];

    /* Message parameters */
    uint8_t msg[SAMPLE_MSG_SIZE];
    uint8_t msg_encrypted[SAMPLE_ENC_MSG_SIZE];
    uint32_t msg_encrypted_len = SAMPLE_ENC_MSG_SIZE;
    uint8_t msg_encrypted_precomputed[SAMPLE_CT_SIZE];
    uint8_t msg_decrypted[SAMPLE_MSG_SIZE];
    uint32_t msg_decrypted_len = SAMPLE_MSG_SIZE;

    printf("AES sample\n");

    /* Prepare data */
    if (hex_string2bytes(IN, msg, SAMPLE_MSG_SIZE)
     || hex_string2bytes(KEY, key, SAMPLE_AES_KEY_SIZE)
     || hex_string2bytes(NONCE, iv, SAMPLE_AES_IV_SIZE)
     || hex_string2bytes(AD, ad, SAMPLE_AES_AD_SIZE)) {
        printf("Error: data preparation failed\n");
        return TEE_ERROR_GENERIC;
    }

    TEE_InitRefAttribute(&attrs[0], TEE_ATTR_SECRET_VALUE, key, SAMPLE_AES_KEY_SIZE);

    if ((rc = TEE_AllocateTransientObject(TEE_TYPE_AES, SAMPLE_AES_MAX_SIZE_BIT, &obj_hndl))
     || (rc = TEE_PopulateTransientObject(obj_hndl, attrs, 1))
     || (rc = TEE_AllocateOperation(&hndl, SAMPLE_AES_ALGO, TEE_MODE_ENCRYPT,
                                    SAMPLE_AES_MAX_SIZE_BIT))
     || (rc = TEE_SetOperationKey(hndl, obj_hndl))) {
        printf("Error: preparation to encryption failed, rc = 0x%08x\n", rc);
        goto L_end;
    }

    if ((rc = TEE_AEInit(hndl, iv, SAMPLE_AES_IV_SIZE, SAMPLE_AES_TAG_SIZE * 8, 0, 0))) {
        printf("Error: AE init for encrypt failed, rc = 0x%08x\n", rc);
        goto L_end;
    }

    TEE_AEUpdateAAD(hndl, ad, SAMPLE_AES_AD_SIZE);

    bytes2hex_string(pbuf, iv, SAMPLE_AES_IV_SIZE);
    printf("IV: %s\n", pbuf);

    printf("Encrypting\n");

    /* TEE_AEUpdate function is used here in vitro / just for example. If final message
     * is prepared, it's reasonable to use TEE_AEEncryptFinal for full message.
     */
    if ((rc = TEE_AEUpdate(hndl, msg, SAMPLE_MSG_INITIAL_LEN, msg_encrypted,
                           &msg_encrypted_len))) {
        printf("Error: AE update failed, rc = 0x%08x\n", rc);
        goto L_end;
    }

    printf("Encrypted %d bytes\n", msg_encrypted_len);
    bytes2hex_string(pbuf, msg_encrypted, msg_encrypted_len);
    printf("CT: %s\n", pbuf);

    msg_encrypted_len = SAMPLE_ENC_MSG_SIZE - SAMPLE_MSG_INITIAL_LEN;
    if ((rc = TEE_AEEncryptFinal(hndl, msg + SAMPLE_MSG_INITIAL_LEN,
                                 SAMPLE_MSG_SIZE - SAMPLE_MSG_INITIAL_LEN,
                                 msg_encrypted + SAMPLE_MSG_INITIAL_LEN,
                                 &msg_encrypted_len, tag, &tag_len))) {
        printf("Error: AE encrypt final failed, rc = 0x%08x\n", rc);
        goto L_end;
    }
    /* Full encrypted message len */
    msg_encrypted_len += SAMPLE_MSG_INITIAL_LEN;

    TEE_FreeOperation(hndl);
    hndl = NULL;
    TEE_CloseObject(obj_hndl);
    obj_hndl = NULL;

    printf("Encrypted final\n");
    bytes2hex_string(pbuf, msg_encrypted, msg_encrypted_len);
    printf("CT: %s\n", pbuf);
    bytes2hex_string(pbuf, tag, tag_len);
    printf("TAG: %s\n", pbuf);

    printf("Decrypting\n");

    if ((rc = TEE_AllocateTransientObject(TEE_TYPE_AES, SAMPLE_AES_MAX_SIZE_BIT, &obj_hndl))
     || (rc = TEE_PopulateTransientObject(obj_hndl, attrs, 1))
     || (rc = TEE_AllocateOperation(&hndl, SAMPLE_AES_ALGO, TEE_MODE_DECRYPT,
                                    SAMPLE_AES_MAX_SIZE_BIT))
     || (rc = TEE_SetOperationKey(hndl, obj_hndl))) {
        printf("Error: preparation to decrypt failed, rc = 0x%08x\n", rc);
        goto L_end;
    }

    if ((rc = TEE_AEInit(hndl, iv, SAMPLE_AES_IV_SIZE, SAMPLE_AES_TAG_SIZE * 8, 0, 0))) {
        printf("Error: AE init for decrypt failed, rc = 0x%08x\n", rc);
        goto L_end;
    }

    TEE_AEUpdateAAD(hndl, ad, SAMPLE_AES_AD_SIZE);

    if ((rc = TEE_AEDecryptFinal(hndl, msg_encrypted, msg_encrypted_len, msg_decrypted,
                                 &msg_decrypted_len, tag, tag_len))) {
        printf("Error: AE decrypt final failed, rc = 0x%08x\n", rc);
        goto L_end;
    }

    printf("Decrypted\n");
    bytes2hex_string(pbuf, msg_decrypted, msg_decrypted_len);
    printf("MSG: %s\n", pbuf);

    printf("Comparing enc/dec results with pre-computed data\n");

    if ((msg_decrypted_len != SAMPLE_MSG_SIZE)
     || (rc = TEE_MemCompare(msg, msg_decrypted, SAMPLE_MSG_SIZE))) {
        bytes2hex_string(pbuf, msg_encrypted_precomputed, SAMPLE_MSG_SIZE);
        printf("Input message differs, rc = 0x%08x\n", rc);
        printf("MSG: %s\n", pbuf);
        goto L_end;
    }
    if (hex_string2bytes(CT, msg_encrypted_precomputed, SAMPLE_CT_SIZE)
     || hex_string2bytes(TAG, tag_precomputed, SAMPLE_AES_TAG_SIZE)) {
        printf("Error: wrong format of data\n");
        rc = TEE_ERROR_GENERIC;
        goto L_end;
    }
    if ((msg_encrypted_len != SAMPLE_CT_SIZE)
     || (rc = TEE_MemCompare(msg_encrypted, msg_encrypted_precomputed, SAMPLE_CT_SIZE))) {
        bytes2hex_string(pbuf, msg_encrypted_precomputed, SAMPLE_CT_SIZE);
        printf("Pre-computed cipher text differs, rc = 0x%08x\n", rc);
        printf("CT: %s\n", pbuf);
        goto L_end;
    }
    if ((tag_len != SAMPLE_AES_TAG_SIZE)
     || (rc = TEE_MemCompare(tag, tag_precomputed, SAMPLE_AES_TAG_SIZE))) {
        bytes2hex_string(pbuf, tag_precomputed, SAMPLE_AES_TAG_SIZE);
        printf("Pre-computed tag differs, rc = 0x%08x\n", rc);
        printf("TAG: %s\n", pbuf);
        goto L_end;
    }

    printf("Enc/dec results and pre-computed data are same\n");

L_end:

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

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

    return rc;
}
