/*
 * tee_dh_key_derive_TA.c
 *
 * DH-2048 key derivation example
 *
 * NOTE: Test data was obtained from TEEGRIS project: tzsl/tests/ta/dh_testvectors.h
 *       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"

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

/* Prime */
static const char P[] = "A2F9418C6CFF9BB5022ED3494227E38E63E9FD9F5171D0733FDE262A508F9E41"
                        "34D297C21A4048E04531A77159C2D4294AF75D299513C913CD23149F7A653896"
                        "7E37C3EBB8B4C6AC46AFFD445A646CA7838EF764E03EA6C3C6CA31B9B7F74E03"
                        "6BA96069FD6D31D8AB128479C356A079D12E7521FA0BF52D03A298BBABA623E9"
                        "D407C5E09A60436EFE99CD7E6500EBAF32C6FC09E45363082C1CA306CE3E9367"
                        "DAE763A2CED4A136B3791BEF38FB4C7503CB7ABAB2401ED3AE9887D2B542685D"
                        "E6916764458DA6BDFA62E785FC0F22D37FAD0DD6C017DB475D11A8FC3E7D319D"
                        "952A1DAE3BFC13C1C75F10BD6BA0BBF188B43B80E88B295196A555B169A47163";
/* Subprime */
static const char Q[] = "";
/* Base */
static const char G[] = "02";

/* DH parameters */
#define SAMPLE_DH_ALGO         TEE_ALG_DH_DERIVE_SHARED_SECRET
#define SAMPLE_DH_SIZE_BIT     2048
#define SAMPLE_DH_SIZE         (SAMPLE_DH_SIZE_BIT / 8)
#define SAMPLE_DH_MAX_SIZE_BIT SAMPLE_DH_SIZE_BIT

#define SAMPLE_P_SIZE          ((sizeof(P) - 1) /2)
#define SAMPLE_Q_SIZE          ((sizeof(Q) - 1) /2)
#define SAMPLE_G_SIZE          ((sizeof(G) - 1) /2)

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

/*
 * Generates DH 2048-bit keys for A and B, generates shared AB secret
 */
TEE_Result dh_key_derive()
{
    TEE_Result rc;
    TEE_Attribute attrs[3];
    size_t nof_attrs = 0;
    TEE_OperationHandle a_handle = NULL;
    TEE_OperationHandle b_handle = NULL;
    TEE_ObjectHandle obj_a_handle = NULL;
    TEE_ObjectHandle obj_b_handle = NULL;
    TEE_ObjectHandle obj_ab_handle = NULL;
    TEE_ObjectHandle obj_ba_handle = NULL;
    /* Print buffer */
    char pbuf[1024] = "\0";
    uint8_t p[SAMPLE_P_SIZE];
    uint8_t q[SAMPLE_Q_SIZE];
    uint8_t g[SAMPLE_G_SIZE];
    uint8_t a_secret_value[SAMPLE_DH_SIZE];
    uint8_t b_secret_value[SAMPLE_DH_SIZE];
    uint8_t a_public_value[SAMPLE_DH_SIZE];
    uint8_t b_public_value[SAMPLE_DH_SIZE];
    uint32_t a_secret_value_len = SAMPLE_DH_SIZE;
    uint32_t b_secret_value_len = SAMPLE_DH_SIZE;
    uint32_t a_public_value_len = SAMPLE_DH_SIZE;
    uint32_t b_public_value_len = SAMPLE_DH_SIZE;

    printf("DH keys derivation sample\n");

    /* Prepare data */
    if (hex_string2bytes(P, p, SAMPLE_P_SIZE)
     || hex_string2bytes(G, g, SAMPLE_G_SIZE)) {
        printf("Error: P,G data preparation failed\n");
        return TEE_ERROR_GENERIC;
    }

    TEE_InitRefAttribute(&attrs[nof_attrs++], TEE_ATTR_DH_PRIME, p, SAMPLE_P_SIZE);
    TEE_InitRefAttribute(&attrs[nof_attrs++], TEE_ATTR_DH_BASE, g, SAMPLE_G_SIZE);
    if (SAMPLE_Q_SIZE != 0) {
        if (hex_string2bytes(Q, q, SAMPLE_Q_SIZE)) {
            printf("Error: Q data preparation failed\n");
            return TEE_ERROR_GENERIC;
        }
        TEE_InitRefAttribute(&attrs[nof_attrs++], TEE_ATTR_DH_SUBPRIME, q, SAMPLE_Q_SIZE);
    }

    printf("Generating keys\n");
    if ((rc = TEE_AllocateTransientObject(TEE_TYPE_DH_KEYPAIR, SAMPLE_DH_MAX_SIZE_BIT,
                                          &obj_a_handle))
     || (rc = TEE_AllocateOperation(&a_handle, SAMPLE_DH_ALGO, TEE_MODE_DERIVE,
                                    SAMPLE_DH_MAX_SIZE_BIT))
     || (rc = TEE_GenerateKey(obj_a_handle, SAMPLE_DH_SIZE_BIT, attrs, nof_attrs))
     || (rc = TEE_GetObjectBufferAttribute(obj_a_handle, TEE_ATTR_DH_PUBLIC_VALUE,
                                          a_public_value, &a_public_value_len))
     || (rc = TEE_SetOperationKey(a_handle, obj_a_handle))) {
        printf("Error: A key generate failed, rc = 0x%08x\n", rc);
        goto L_end;
    }

    if ((rc = TEE_AllocateTransientObject(TEE_TYPE_DH_KEYPAIR, SAMPLE_DH_MAX_SIZE_BIT,
                                          &obj_b_handle))
     || (rc = TEE_AllocateOperation(&b_handle, SAMPLE_DH_ALGO, TEE_MODE_DERIVE,
                                    SAMPLE_DH_MAX_SIZE_BIT))
     || (rc = TEE_GenerateKey(obj_b_handle, SAMPLE_DH_SIZE_BIT, attrs, nof_attrs))
     || (rc = TEE_GetObjectBufferAttribute(obj_b_handle, TEE_ATTR_DH_PUBLIC_VALUE,
                                          b_public_value, &b_public_value_len))
     || (rc = TEE_SetOperationKey(b_handle, obj_b_handle))) {
        printf("Error: B key generate failed, rc = 0x%08x\n", rc);
        goto L_end;
    }

    if ((rc = TEE_AllocateTransientObject(TEE_TYPE_GENERIC_SECRET, SAMPLE_DH_MAX_SIZE_BIT,
                                          &obj_ab_handle))
     || (rc = TEE_AllocateTransientObject(TEE_TYPE_GENERIC_SECRET, SAMPLE_DH_MAX_SIZE_BIT,
                                          &obj_ba_handle))) {
        printf("Error: DH allocate generic secret failed, rc = 0x%08x\n", rc);
        goto L_end;
    }

    printf("Deriving keys\n");

    TEE_InitRefAttribute(&attrs[0], TEE_ATTR_DH_PUBLIC_VALUE, b_public_value, b_public_value_len);
    TEE_DeriveKey(a_handle, attrs, 1, obj_ab_handle);

    TEE_InitRefAttribute(&attrs[0], TEE_ATTR_DH_PUBLIC_VALUE, a_public_value, a_public_value_len);
    TEE_DeriveKey(b_handle, attrs, 1, obj_ba_handle);

    printf("Comparing secret values A and B\n");

    if ((rc = TEE_GetObjectBufferAttribute(obj_ab_handle, TEE_ATTR_SECRET_VALUE, a_secret_value,
                                           &a_secret_value_len))
     || (rc = TEE_GetObjectBufferAttribute(obj_ba_handle, TEE_ATTR_SECRET_VALUE, b_secret_value,
                                           &b_secret_value_len))) {
        printf("Error: get secret value failed, rc = 0x%08x\n", rc);
        goto L_end;
    }

    bytes2hex_string(pbuf, a_secret_value, a_secret_value_len);
    printf("Secret A, %d bit: %s\n", a_secret_value_len << 3, pbuf);

    if ((a_secret_value_len != b_secret_value_len)
     || (rc = TEE_MemCompare(a_secret_value, b_secret_value, a_secret_value_len))) {
         printf("Error: Secret B differs, rc = 0x%08x\n", rc);
         bytes2hex_string(pbuf, b_secret_value, b_secret_value_len);
         printf("Secret B, %d bit: %s\n", b_secret_value_len << 3, pbuf);
         goto L_end;
    }

    printf("Secret B is the same\n");

L_end:

    if (a_handle) {
        TEE_FreeOperation(a_handle);
    }
    if (b_handle) {
        TEE_FreeOperation(b_handle);
    }

    TEE_CloseObject(obj_a_handle);
    TEE_CloseObject(obj_b_handle);
    TEE_CloseObject(obj_ab_handle);
    TEE_CloseObject(obj_ba_handle);

    if (rc) {
        printf("DH derive failed\n");
    } else {
        printf("DH derive passed\n");
    }

    return rc;
}
