/*
 * tee_crypto_samples_CA.c
 *
 * Crypto samples client application (NWd part)
 *
 */

#include "ta/crypto_samples_commandid.h"
#include "tee_client_api.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

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

#define _EXIT(res) exit(res)
#define _EXIT_P(res) \
    do {\
        if (res) {\
            printf("Error: something went wrong, line %d\n", __LINE__);\
        }\
        exit(res);\
    } while (0)

#define SUCCESS TEEC_SUCCESS  /* 0 */
#define FAIL    TEEC_IMP_MIN  /* 1 */

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

TEEC_UUID uuid = {
    .timeLow = 0,
    .timeMid = 0,
    .timeHiAndVersion = 0,
    .clockSeqAndNode = {0, 0, 0, 0, 0, 0, 0, 0x67},
};

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

static const char *Sample_name[NUMBER_OF_SAMPLES] = {
    [ATTRIBUTES_GET_SET_SAMPLE] = "Attributes get/set",
    [SHA256_DIGEST_SAMPLE] = "SHA-256",
    [HMAC_SHA256_DIGEST_SAMPLE] = "HMAC-SHA-256",
    [ECDSA_SIGN_VERIFY_SAMPLE] = "ECDSA sign/verify",
    [RSA_SIGN_VERIFY_SAMPLE] = "RSA sign/verify",
    [RSA_ENC_DEC_SAMPLE] = "RSA encrypt/decrypt",
    [AES_ENC_DEC_SAMPLE] = "AES GCM encrypt/decrypt",
    [DH_KEY_DERIVE_SAMPLE] = "DH key derive",
};

#define NUMBER_OF_ID    NUMBER_OF_SAMPLES

typedef struct {
    uint8_t run_mode;
    uint8_t fail_mode;
    size_t id_list_len;
    unsigned id_list[NUMBER_OF_ID];
} CLI_Options;

enum {
    FAIL_MODE_CONTINUE = 0,
    FAIL_MODE_STOP,
};

enum {
    RUN_MODE_ALL = 0,
    RUN_MODE_SELECTED,
};

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

static void print_usage()
{
    printf("\nCrypto samples application\n");
    printf("Usage: ./tee_crypto_samples [options [<args>]]\n");
    printf("Options:\n");
    printf("    -h           print this help and exit\n");
    printf("    -f           stop run if any of crypto samples fails\n");
    printf("    -s <args>    run crypto samples in <args> list\n");
    for (unsigned i = 0; i < NUMBER_OF_SAMPLES; i++) {
        printf("        %u - %s\n", i, Sample_name[i]);
    }
    printf("Usage examples:\n");
    printf("    ./tee_crypto_samples\n");
    printf("    ./tee_crypto_samples -s 0 1 2 3\n");
    printf("    ./tee_crypto_samples -fs 0 6 5\n");
    printf("\n" );
}

static void process_cli_arguments(int argc, char *argv[], CLI_Options *options)
{
    int c;
    unsigned id;

    if (!options) {
        _EXIT_P(FAIL);
    }
    memset(options, 0, sizeof(CLI_Options));

    while ((c = getopt(argc, argv, "fs:h")) != -1) {
        switch (c) {
            case 's':
                options->run_mode = RUN_MODE_SELECTED;
                for (int i = optind - 1; i < argc; i++) {
                    if (!sscanf(argv[i], "%u", &id)) {
                        break;
                    }
                    if (options->id_list_len >= NUMBER_OF_ID) {
                        printf("Warning: samples id-list is full\n");
                        break;
                    }
                    if (id >= NUMBER_OF_SAMPLES) {
                        printf("Error: unknown sample id = %d\n", id);
                        _EXIT(FAIL);
                    }
                    options->id_list[options->id_list_len] = id;
                    options->id_list_len++;
                }
                if (!options->id_list_len) {
                    printf("Error: no selected samples\n");
                    _EXIT(FAIL);
                }
                break;
            case 'f':
                options->fail_mode = FAIL_MODE_STOP;
                break;
            case 'h':
                print_usage();
                _EXIT(SUCCESS);
            default:
                print_usage();
                _EXIT(FAIL);
        }
    }

    if (options->run_mode == RUN_MODE_ALL) {
        for (unsigned i = 0; i < NUMBER_OF_SAMPLES; i++) {
            options->id_list[i] = i;
        }
        options->id_list_len = NUMBER_OF_SAMPLES;
    }

}

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

/*
 * Runs all or selected from command line crypto samples
 */
int main(int argc, char* argv[])
{
    TEEC_Result res = FAIL;
    TEEC_Context context;
    TEEC_Session session;
    uint32_t return_origin;
    CLI_Options opt;

    process_cli_arguments(argc, argv, &opt);

    printf("Initialize context\n");
    if ((res = TEEC_InitializeContext(NULL, &context))) {
        printf("Error: TEEC_InitializeContext failed, res = 0x%08x\n", res);
        return SUCCESS;
    }

    printf("Open session\n");
    if ((res = TEEC_OpenSession(&context, &session, &uuid, 0, NULL, NULL, &return_origin))) {
        printf("Error: TEEC_OpenSession failed, res = 0x08%x\n", res);
        TEEC_FinalizeContext(&context);
        return res;
    }

    printf("Run\n");
    for (unsigned i = 0, id; i < opt.id_list_len; i++) {
        id = opt.id_list[i];
        printf("%u.%s\n", id, Sample_name[id]);
        if ((res = TEEC_InvokeCommand(&session, id, NULL, &return_origin))
         || (return_origin != TEEC_ORIGIN_TRUSTED_APP)) {
            printf("%u.%s failed, res = 0x%08x\n", id, Sample_name[id], res);
            if (opt.fail_mode) {
                return FAIL;
            }
        } else {
            printf("%u.%s successfully executed\n", id, Sample_name[id]);
        }
    }

    TEEC_CloseSession(&session);
    TEEC_FinalizeContext(&context);

    printf("Done\n");

    return res;
}
