/**
 * @file   TigerTests.c
 * @brief  Tiger TA Tests
 * @author Andrey Orlenko (a.orlenko@samsung.com)
 * @date   Created Jul 8, 2016
 *
 * Copyright (c) 2016 Samsung Electronics Co., Ltd. All rights reserved.
 *
 * Created in Samsung Ukraine R&D Center (SRK) under a contract between
 * LLC "Samsung Electronics Ukraine Company" (Kiev, Ukraine)
 * and "Samsung Electronics Co", Ltd (Seoul, Republic of Korea)
 **/

#include "TigerTests.h"
#include "TigerLogging.h"
#include "TigerCore.h"
#include "TigerMacroses.h"
#include "TigerMbedTlsExt.h"
#include "TigerTeeApiWrapper.h"
#include "TigerSignMessage.h"
#include "TigerAttk.h"
#include "TigerSskds.h"
#include "TigerKeyDataStore.h"
#include "TigerCounterDataStore.h"
#include "TigerFingerprintIdTable.h"
#include "tzWrappers/TzwRpmb.h"
#include "tzWrappers/TzwSerialNumber.h"
#include <inttypes.h>

#ifdef TZ_MODEL_QCOM
    #include <qsee_rsa.h>
    #include <qsee_timer.h>
#endif

#include <stdio.h>

#define AUTHHAT_FID_SIZE 32 // TODO (k.volobuyev): get size from bio buffer header

const int KEY_EXPONENT = 65537;

static TEE_Result checkArbitrarySaltLength_Verified();
static TEE_Result exportAuthKey(TigerSessionContext_t* apCtx, void* apRsp);
static TEE_Result genKeyPerformance();
static TEE_Result rpmbReadWriteSuccess();
static TEE_Result incrementCounter();
static TEE_Result checkFingerprintIdTable();
static TEE_Result bioResult();

TEE_Result runTeeApiTests(TigerSessionContext_t* apCtx,
                          const void* apReq, const uint32_t sizeRequest,
                          void* apRsp, const uint32_t sizeResponse) {
    LOG_FUNC_BEGIN;

    S_VAR_NOT_USED(apCtx);
    S_VAR_NOT_USED(apReq);
    S_VAR_NOT_USED(sizeRequest);
    S_VAR_NOT_USED(apRsp);
    S_VAR_NOT_USED(sizeResponse);

    TEE_Result status = TEE_ERROR_NOT_IMPLEMENTED;
    if ((apReq) && (sizeof(uint32_t) == sizeRequest)) {
        uint32_t cmd = 0;
        tzwMemMove(&cmd, apReq, sizeRequest);
        switch (cmd) {
            case TIGER_TEST_TCI_CMD_ID_checkArbitrarySaltLength_Verified:
                status = checkArbitrarySaltLength_Verified();
                break;
            case TIGER_TEST_TCI_CMD_ID_exportAuthKey:
                status = exportAuthKey(apCtx, apRsp);
                break;
            case TIGER_TEST_TCI_CMD_ID_genKeyPerformance:
                status = genKeyPerformance();
                break;
            case TIGER_TEST_TCI_CMD_ID_checkRpmbReadWrite:
                status = rpmbReadWriteSuccess();
                break;
            case TIGER_TEST_TCI_CMD_ID_incrementCounter:
                status = incrementCounter();
                break;
            case TIGER_TEST_TCI_CMD_ID_checkFingerprintIdTable:
                status = checkFingerprintIdTable();
                break;
            case TIGER_TEST_TCI_CMD_ID_checkBioResult:
                status = bioResult();
                break;
            default:
                status = TEE_ERROR_NOT_IMPLEMENTED;
                break;
        }
    }

    LOG_D("status = 0x%08x", status);

    LOG_FUNC_END;
    return status;
}


TEE_Result checkArbitrarySaltLength_Verified() {
    LOG_FUNC_BEGIN;

    mbedtls_rsa_context rsaCtx;
    unsigned char hash[32];
    const int saltLen = 20; // For 32 bytes it is ok
    unsigned char sig[MBEDTLS_MPI_MAX_SIZE];

    int ret = 0;
    TEE_Result status = TEE_SUCCESS;
    do {
        tigerMbedtlsRng(NULL, hash, sizeof(hash));
        logByteArrayBase64(hash, sizeof(hash), "HASH");

        mbedtls_rsa_init( &rsaCtx, MBEDTLS_RSA_PKCS_V21, MBEDTLS_MD_SHA256);

        ret = mbedtls_rsa_gen_key(&rsaCtx, tigerMbedtlsRng, NULL, 2048, KEY_EXPONENT);
        TIGER_CHECK_MBEDTLS_RET_ZERO_BREAK("mbedtls_rsa_gen_key");
        LOG_I("RSA key successfully generated");

        ret = tigerMbedtlsRsaRsassaPssSign(&rsaCtx, tigerMbedtlsRng, NULL, MBEDTLS_RSA_PRIVATE, rsaCtx.hash_id,
                0, hash, saltLen, sig);
        TIGER_CHECK_MBEDTLS_RET_ZERO_BREAK("tigerMbedtlsRsaRsassaPssSign");
        LOG_I("Specialized signature estimated. Salt length: %d bytes", saltLen);
        logByteArrayBase64(sig, rsaCtx.len, "SIGNATURE");

        ret = mbedtls_rsa_rsassa_pss_verify(&rsaCtx, tigerMbedtlsRng, NULL, MBEDTLS_RSA_PUBLIC,
                rsaCtx.hash_id, 0, hash, sig);
        TIGER_CHECK_MBEDTLS_RET_ZERO_BREAK("mbedtls_rsa_rsassa_pss_verify");

        LOG_I("Signature verification success!");

    } while (0);

    mbedtls_rsa_free(&rsaCtx);
    if (ret < 0) {
        status = TEE_ERROR_GENERIC;
    }

    LOG_FUNC_END;

    return status;
}

TEE_Result exportAuthKey(TigerSessionContext_t* apCtx, void* apRsp) {

    TIGER_CHECK_FUNCTION_ARGUMENT_RETURN(apCtx != NULL);
    TIGER_CHECK_FUNCTION_ARGUMENT_RETURN(apCtx->data != NULL);
    TIGER_CHECK_FUNCTION_ARGUMENT_RETURN(apRsp != NULL);
    TIGER_CHECK_FUNCTION_ARGUMENT_RETURN(TCI_TEE_SESSION_TYPE_SIGN == apCtx->type);


    return TEE_SUCCESS;
}
#ifdef TZ_MODEL_QCOM
TEE_Result genKeyPerformance() {
    LOG_FUNC_BEGIN;

    TEE_Result status = TEE_SUCCESS;

#ifdef TZ_MODEL_QCOM
    mbedtls_pk_context ctx;
    QSEE_RSA_KEY* key = NULL;

    unsigned char pubExp[] = {0x00, 0x00, 0x01, 0x00, 0x01};
    const uint32_t NUM_OF_REPEATS = 10;

    do {
        // mbedtls keypair generation

        mbedtls_pk_init(&ctx);

        int ret = mbedtls_pk_setup(&ctx, mbedtls_pk_info_from_type(MBEDTLS_PK_RSA));
        if (ret) {
            logMbedtlsError(ret, __func__);
            status = TEE_ERROR_GENERIC;
            break;
        }

        unsigned long long totalTime = qsee_get_uptime();

        for (uint32_t i = 0; i < NUM_OF_REPEATS; ++i) {
            ret = mbedtls_rsa_gen_key(mbedtls_pk_rsa(ctx), tigerMbedtlsRng, NULL, 2048, KEY_EXPONENT);
            if (ret) {
                logMbedtlsError(ret, __func__);
                status = TEE_ERROR_GENERIC;
                break;
            }
        }

        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("mbedtls_rsa_gen_key");

        totalTime = qsee_get_uptime() - totalTime;
        LOG_I("mbedtls total time: %u ms", totalTime);

        // QSEE keypair generation

        key = (QSEE_RSA_KEY*) tzwMalloc(sizeof(QSEE_RSA_KEY));
        if (NULL == key) {
            LOG_E("QSEE RSA key allocation failed.");
            status = TEE_ERROR_GENERIC;
            break;
        }

        if (qsee_util_init_s_bigint(&(key->e)) || qsee_util_init_s_bigint(&(key->d)) ||
            qsee_util_init_s_bigint(&(key->N)) || qsee_util_init_s_bigint(&(key->p)) ||
            qsee_util_init_s_bigint(&(key->q)) || qsee_util_init_s_bigint(&(key->qP)) ||
            qsee_util_init_s_bigint(&(key->dP)) || qsee_util_init_s_bigint(&(key->dQ))) {
            LOG_E("Allocate and initialize S_BIGINT data failed");
            status = TEE_ERROR_GENERIC;
            break;
        }

        totalTime = qsee_get_uptime();

        for (uint32_t i = 0; i < NUM_OF_REPEATS; ++i) {
            int res = qsee_rsa_key_gen(key, 256, pubExp, sizeof(pubExp));
            if (res) {
                LOG_E("qsee_rsa_key_gen failed with return code: %d\n", res);
                status = TEE_ERROR_GENERIC;
                break;
            }
        }

        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("qsee_rsa_key_gen");

        totalTime = qsee_get_uptime() - totalTime;
        LOG_I("QSEE total_time: %u ms", totalTime);
    } while (0);

    mbedtls_pk_free(&ctx);
    tzwFree(key);
#endif

    // TODO(k.volobuyev): Add similar test for Blowfish

    LOG_FUNC_END;
    return status;
}
#else
TEE_Result genKeyPerformance() {
    return TEE_SUCCESS;
}
#endif

TEE_Result rpmbReadWriteSuccess() {
    LOG_FUNC_BEGIN;
    const uint32_t randomBytesLen = 1124; //Rsa key-pair 2048 der size
    uint8_t bufferToWrite[MAX_EXPORT_SIZE_BYTES] = {0};
    uint8_t bufferToRead[MAX_EXPORT_SIZE_BYTES] = {0};

    tzwGenerateRandom(bufferToWrite, randomBytesLen);

    TEE_Result status = rpmbWrite(bufferToWrite, randomBytesLen);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("rpmbWrite");

    status = rpmbRead(bufferToRead, sizeof(bufferToRead));
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("rpmbRead");

    int32_t ret = tzwMemCompare((void*) bufferToWrite, (void*) bufferToRead, randomBytesLen);
    if (0 != ret) {
        LOG_E("Data did not match!");
        return TEE_ERROR_CORRUPT_OBJECT;
    }
    return status;
}

TEE_Result incrementCounter() {
    TigerCounter_t counterInit = 0;
    TEE_Result status = tigerGetCounter(&counterInit);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("tigerGetCounter");

    status = tigerIncCounter();
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("tigerIncCounter");

    TigerCounter_t counterFin = 0;
    status = tigerGetCounter(&counterFin);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("tigerGetCounter");

    if ((counterInit + 1) != counterFin) {
        LOG_E("Counter did not increment!");
        status = TEE_ERROR_BAD_STATE;
    }

    return status;
}


TEE_Result checkFingerprintIdTable() {
    uint32_t soterFid = 3;
    uint32_t realFidSize = AUTHHAT_FID_SIZE;
    uint64_t entryIndex = 0;
    uint8_t realFid[AUTHHAT_FID_SIZE];
    TEE_Result status = TEE_SUCCESS;

    // check the table with real FID which was got from the authhat

    // FIXME (k.volobuyev): Replace this with calling of function which get real FID from the authhat
    tzwMemFill(realFid, 42, AUTHHAT_FID_SIZE);

    status = tigerFPTableInsert(realFid, realFidSize, entryIndex, soterFid);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("tigerFPTableAdd");

    status = tigerFPTableGetUniqueId(realFid, realFidSize, entryIndex, &soterFid);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("tigerFPTableGetUniqueId");

    LOG_D("SOTER FID based on real FID is: %u", soterFid);

    status = tigerFPTableSave();
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("tigerFPTableSave");

    // rewrite already existed real FID

    const uint32_t MAX_ENROLLED_FID_NUM = 10;

    // taking into account that numeration of SOTER FID begins from 3
    const uint32_t EXPECTED_SOTER_FID = MAX_ENROLLED_FID_NUM + 3;

    // begin from second cell to skip the real FID which was got from the authhat.
    // + 1 was added to rewrite already existed real FID
    for (uint32_t i = 1; i < MAX_ENROLLED_FID_NUM + 1; i++) {
        tzwMemFill(realFid, i, AUTHHAT_FID_SIZE);
        entryIndex = i;

        // use FID index with already existed real FID
        if (MAX_ENROLLED_FID_NUM == i) {
            entryIndex = 1;
        }

        status = tigerFPTableLoad();
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("tigerFPTableLoad");

        status = tigerFPTableGetUniqueId(realFid, realFidSize, entryIndex, &soterFid);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("tigerFPTableGetUniqueId");

        LOG_D("Current SOTER FID is: %u", soterFid);

        // check if new SOTER FID was generated for newly rewritten real FID
        if (MAX_ENROLLED_FID_NUM == i) {
            TIGER_CHECK_FUNCTION_ARGUMENT_RETURN(soterFid == EXPECTED_SOTER_FID);
        }

        status = tigerFPTableSave();
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("tigerFPTableSave");
    }
    return status;
}

static TEE_Result bioResult() {
    const uint32_t badFPIndex = 128;
    uint32_t fidBufferSize = 128;
    uint32_t fidExpectedSize = 32;

    uint8_t* fid = tzwMalloc(fidBufferSize);
    uint32_t fpIndex = badFPIndex;
    uint64_t challenge = 0;
    TEE_Result status = TEE_ERROR_GENERIC;
    do {
        status = getLatestAuthResult(fid, &fidBufferSize, &fpIndex);
        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("getLatestAuthResult");
        TIGER_CHECK_CONDITION_BREAK(fpIndex != badFPIndex);
        TIGER_CHECK_CONDITION_BREAK(challenge != 0);
        LOG_D("fidBufferSize = %u, fidExpectedSize = %u", fidBufferSize, fidExpectedSize);
        TIGER_CHECK_CONDITION_BREAK(fidBufferSize == fidExpectedSize);
        status = TEE_ERROR_CORRUPT_OBJECT;
        // Check that fid not zero.
        for (size_t i = 0; i < fidBufferSize; ++i) {
            if (fid[i] != 0) {
                status = TEE_SUCCESS;
                break;
            }
        }
    } while(0);

    return status;
}

TEE_Result fillTestDeviceId(uint8_t* deviceId, const uint32_t devIdSize) {
    LOG_FUNC_BEGIN;

    TIGER_CHECK_FUNCTION_ARGUMENT_RETURN(deviceId != NULL);
    TIGER_CHECK_FUNCTION_ARGUMENT_RETURN(devIdSize == DEVICE_ID_SIZE);

    static const char* testIdPrefix = "SRKTEST_";
    const size_t testIdPrefixLen = strlen(testIdPrefix);
    const size_t randomBytesLen = (devIdSize - testIdPrefixLen) / BYTE_IN_HEX_LEN; // because will be converted to hex

    uint8_t* randomBytes = NULL;

    TEE_Result status = TEE_SUCCESS;
    do {
        randomBytes = tzwMalloc(randomBytesLen);
        TIGER_CHECK_BUFFER_ALLOCATED_BREAK(randomBytes);

        int ret = tzwGenerateRandom(randomBytes, randomBytesLen);
        if (0 != ret) {
            status = TEE_ERROR_GENERIC;
            break;
        }

        // copy prefix
        tzwMemMove(deviceId, testIdPrefix, testIdPrefixLen);

        convertBytesToHex(randomBytes, randomBytesLen, (deviceId + testIdPrefixLen), randomBytesLen * BYTE_IN_HEX_LEN);
        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("convertBytesToHex");
    } while(0);

    tzwFree(randomBytes);

    LOG_FUNC_END;
    return status;
}

TEE_Result getTestAttk(TigerKeyPair_t* keyPair) {
    LOG_FUNC_BEGIN;
    TigerObjectId_t* objId = NULL;
    const char* charAlias = "ATTK_TEST_OBJ";
    const TciProcessUid_t uidSelf = 1;

    TciKeyAlias_t keyAlias = {0, {0}};
    tzwMemMove(keyAlias.value, charAlias, strlen(charAlias));
    keyAlias.length = strlen(charAlias);

    TEE_Result status = tigerCreateKeyPairObjectId(&keyAlias, uidSelf, &objId);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("tigerCreateKeyPairObjectId");

    LOG_D("Fake attk object id(%s) generated!", charAlias);
    status = tigerLoadKeyPair(objId, keyPair);

    do {
        if (TEE_SUCCESS != status) {
            LOG_D("Test attk not exist, generate new one.");
            status = tigerGenerateKeyPair(keyPair);
            TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("tigerGenerateKeyPair");

            LOG_D("Test attk generated!, save it!");

            status = tigerSaveKeyPair(objId, keyPair);


            TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("tigerSaveKeyPair");

            LOG_D("Test attk saved!");
        }
    } while(0);

    tigerFreeObjectId(objId);
    LOG_FUNC_END;
    return status;
}

TEE_Result getTestAttkPubBytes(uint8_t* buf, uint32_t* buflen) {
    LOG_FUNC_BEGIN;

    TigerKeyPair_t* kp = tigerAllocateKeyPair();
    TIGER_CHECK_BUFFER_ALLOCATED_RETURN(kp);

    TEE_Result status = getTestAttk(kp);

    if (status == TEE_SUCCESS) {
        LOG_D("Test ATTK was loaded successfully., size(%u)", *buflen);

        int ret = mbedtls_pk_write_pubkey_pem(tigerGetKeyPairContext(kp), (unsigned char*)buf, *buflen);

        LOG_D("Printing ATTK.");
        if (0 != ret) {
            logMbedtlsError(ret, "Could not write attk pem.");
            status = TEE_ERROR_CORRUPT_OBJECT;
        } else {
            logByteArrayString(buf, strlen((const char*) buf), "ATTK PUB");
        }
    }

    tigerFreeKeyPair(kp);

    LOG_FUNC_END;
    return status;
}