/*
 * 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)
 */

/**
 * @file TigerDataKeyStore.c
 * @brief Implements functionality to replace Global Platform TEE_TransientObject.
 * @author Viktor Kopp (v.kopp@samsung.com)
 * @date Created Jul 15, 2016
 */

#include "TigerKeyDataStore.h"

#include <assert.h>
#include <stdio.h>
#include <string.h>

#include "mbedtls/asn1.h"
#include "mbedtls/asn1write.h"
#include "mbedtls/pk.h"

#include "TigerAttk.h"
#include "TigerCounterDataStore.h"
#include "TigerJson.h"
#include "TigerLogging.h"
#include "TigerMacroses.h"
#include "TigerMbedTlsExt.h"
#include "TigerPemUtils.h"
#include "TigerSskds.h"
#include "TzwRpmb.h"
#include "TzwMemory.h"
#include "TzwString.h"

#define TIGER_DEFAULT_KEY_BUFFER_SIZE 2048
#define TIGER_DEFAULT_KEY_SIZE 2048
#define TIGER_TAG_AND_LEN_BUFFER_SIZE 8
#define TIGER_MOVE_BUFFER_SIZE 512
#define TIGER_CERT_SIGN_SALT_LENGTH  20 /*bytes*/
#define TIGER_CERT_SIGN_LENGTH      256

#define MAX_NB_DIGITS_IN_UINT32 10
#define MAX_NB_DIGITS_IN_UINT64 20

#define MBEDTLS_HASH_SIZE 64
// This value was got from mbedtls-2.2.1/library/pkwrite.c
#define RSA_PUB_DER_MAX_BYTES 38 + 2 * MBEDTLS_MPI_MAX_SIZE

struct TigerKeyPair {
    mbedtls_pk_context ctx;
};

typedef enum TigerKeyPart {
    TKP_KEY_PAIR = 0,
    TKP_KEY_INFO = 1,
    TKP_KEY_CHLD = 2,
} TigerKeyPart_t;

static int serializeKeyPair(const TigerKeyPair_t* const keypair, uint8_t* buf, size_t len);
static TEE_Result seekToKeyPair(TzwSfsObject_t obj, size_t* len);
static TEE_Result seekTo(TzwSfsObject_t obj, TigerKeyPart_t part, size_t* len);
static TEE_Result writeJsonCertificate(JsonObject_t jsonCertObj, const uint8_t* pemNewLinesReplaced,
                                const uint32_t pemNewLinesReplacedSize,
                                const uint8_t* cpuId, const uint32_t cpuIdSize,
                                const uint64_t counter, uint32_t uid);

TEE_Result tigerSaveKeyPair(const TigerObjectId_t* const objId,
        const TigerKeyPair_t* const keypair) {
    LOG_FUNC_BEGIN;

    TIGER_ASSERT(NULL != objId);
    TIGER_ASSERT(NULL != keypair);

    uint8_t* buf = tzwMalloc(TIGER_DEFAULT_KEY_BUFFER_SIZE);
    TIGER_CHECK_BUFFER_ALLOCATED_RETURN(buf);

    const uint8_t* buf_end = buf + TIGER_DEFAULT_KEY_BUFFER_SIZE;

    TEE_Result status = TEE_SUCCESS;
    TzwSfsObject_t handle = TEE_HANDLE_NULL;
    int bytes_serialized = 0;
    do {
        // "on successful open-call delete" is used for compatibility with GP, where direct TEE_CreatePersistentObject
        // did not recognize flag TEE_DATA_FLAG_OVERWRITE
        status = tzwOpenSfsObject((void*) tigerGetObjectIdData(objId), tigerGetObjectIdLength(objId),
                TEE_DATA_FLAG_ACCESS_WRITE | TEE_DATA_FLAG_ACCESS_WRITE_META, &handle);
        if (TEE_SUCCESS == status) {
            LOG_D("Object existed. Deleting.");
            tzwCloseAndDeleteSfsObject(handle);
        }

        status = tzwCreateSfsObject((void*) tigerGetObjectIdData(objId), tigerGetObjectIdLength(objId),
                                    TEE_DATA_FLAG_ACCESS_WRITE | TEE_DATA_FLAG_ACCESS_WRITE_META,
                                    &handle);
        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("tzwCreateSfsObject");

        // serialize key pair data.
        bytes_serialized = serializeKeyPair(keypair, buf, TIGER_DEFAULT_KEY_BUFFER_SIZE);
        if (bytes_serialized < 0) {
            status = TEE_ERROR_GENERIC;
            break;
        }
        LOG_D("Key pair serialized in %d bytes", bytes_serialized);

        // flush data to persistent object
        status = tzwWriteToSfsObject(handle, (void*) (buf_end - bytes_serialized), bytes_serialized);
        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("tzwWriteToSfsObject");

        LOG_I("Data is written successfully!");
    } while (0);

    if (bytes_serialized < 0) {
        logMbedtlsError(bytes_serialized, __func__);
    }

    tzwCloseSfsObject(handle);
    tzwFree(buf);

    LOG_FUNC_END;
    return status;
}

TEE_Result tigerLoadKeyPair(const TigerObjectId_t* const objId, TigerKeyPair_t* result) {
    LOG_FUNC_BEGIN;

    TIGER_ASSERT(NULL != objId);
    TIGER_ASSERT(NULL != result);

    logByteArrayHex(tigerGetObjectIdData(objId), tigerGetObjectIdLength(objId), "OBJ_ID");

    TzwSfsObject_t handle = TEE_HANDLE_NULL;
    TEE_Result status = TEE_SUCCESS;
    uint8_t* buf = NULL;
    do {
        status = tzwOpenSfsObject((void*) tigerGetObjectIdData(objId), tigerGetObjectIdLength(objId),
                TEE_DATA_FLAG_ACCESS_READ, &handle);
        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("tzwOpenSfsObject");
        logObject(handle);

        size_t len = 0;
        status = seekToKeyPair(handle, &len);
        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("seekToKeyPair");
        logObject(handle);

        buf = tzwMalloc(TIGER_DEFAULT_KEY_BUFFER_SIZE);
        TIGER_CHECK_BUFFER_ALLOCATED_BREAK(buf);

        uint32_t count = 0;
        tzwReadFromSfsObject(handle, buf, len, &count);
        if (count != len) {
            LOG_E("tzwReadFromSfsObject() failed to read keypair data.");
            status = TEE_ERROR_GENERIC;
            break;
        }

        int ret = 0;
        //mbedtls_pk_init(&result->ctx);
        ret = mbedtls_pk_parse_key(&result->ctx, buf, len, NULL, 0);
        if (0 != ret) {
            logMbedtlsError(ret, __func__);
            status = TEE_ERROR_GENERIC;
            break;
        }
        LOG_D("Pk parsed successfully!");
    } while (0);

    tzwFree(buf);
    tzwCloseSfsObject(handle);

    LOG_FUNC_END;
    return status;
}

void* tigerGetKeyPairContext(const TigerKeyPair_t* kp) {
    assert(kp != NULL);

    return (void*) &kp->ctx;
}

TEE_Result tigerDeleteKeyPair(const TigerObjectId_t* const objId) {
    LOG_FUNC_BEGIN;
    TIGER_ASSERT(objId != NULL);

    TEE_Result status = TEE_SUCCESS;
    TzwSfsObject_t handle = TEE_HANDLE_NULL;
    do {
        status = tzwOpenSfsObject((void*) tigerGetObjectIdData(objId), tigerGetObjectIdLength(objId),
                TEE_DATA_FLAG_ACCESS_WRITE | TEE_DATA_FLAG_ACCESS_WRITE_META,
                &handle);
        if ((TEE_SUCCESS != status) && (TEE_ERROR_ITEM_NOT_FOUND != status)) {
            logTeeError(status, __func__);
        }
    } while (0);

    tzwCloseAndDeleteSfsObject(handle);

    LOG_FUNC_END;
    return status;
}

TEE_Result tigerGenerateKeyPair(TigerKeyPair_t* keypair) {
    TIGER_ASSERT(keypair != NULL)

    int ret = 0;
    TEE_Result status = TEE_SUCCESS;
    do {
        ret = mbedtls_pk_setup(&keypair->ctx, mbedtls_pk_info_from_type(MBEDTLS_PK_RSA));
        if (0 != ret) {
            status = TEE_ERROR_GENERIC;
            break;
        }

        const int defaultRsaPublicExponent = 65537;
        ret = mbedtls_rsa_gen_key(mbedtls_pk_rsa(keypair->ctx), tigerMbedtlsRng, NULL, TIGER_DEFAULT_KEY_SIZE, defaultRsaPublicExponent);
        if (0 != ret) {
            status = TEE_ERROR_GENERIC;
            break;
        }
    } while (0);

    if (0 != ret) {
        logMbedtlsError(ret, __func__);
    }
    return status;
}

TEE_Result tigerGetPublicKey(const TigerObjectId_t* const objId, uint8_t** p, size_t size) {
    TigerKeyPair_t* kp = tigerAllocateKeyPair();
    TIGER_CHECK_BUFFER_ALLOCATED_RETURN(kp);

    TEE_Result status = TEE_SUCCESS;
    do {
        status = tigerLoadKeyPair(objId, kp);
        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("tigerLoadKeyPair");

        int ret = mbedtls_pk_write_pubkey_der(&kp->ctx, *p - size, size);
        if (ret <= 0) {
            logMbedtlsError(ret, __func__);
            status = TEE_ERROR_GENERIC;
            break;
        }
        *p -= ret;
    } while (0);
    tigerFreeKeyPair(kp);

    return status;
}

TEE_Result tigerGenerateJsonCertificate(const TigerObjectId_t* objId, const TigerObjectId_t* parentObjId, TciProcessUid_t uid, uint8_t* p, uint32_t* size) {
    LOG_FUNC_BEGIN;

    TIGER_ASSERT(p != NULL)
    TIGER_ASSERT(objId != NULL)
    TIGER_ASSERT(size != NULL)

    TigerKeyPair_t* kp = NULL;
    TigerKeyPair_t* kpParent = NULL;
    JsonObject_t jsonCertObj = NULL;
    uint8_t* pemNewLinesReplaced = NULL;

    int ret = 0;
    TEE_Result status = TEE_SUCCESS;
    do {
        kp = tigerAllocateKeyPair();
        TIGER_CHECK_BUFFER_ALLOCATED_BREAK(kp);

        status = tigerLoadKeyPair(objId, kp);
        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("tigerLoadKeyPair");

        kpParent = tigerAllocateKeyPair();
        TIGER_CHECK_BUFFER_ALLOCATED_BREAK(kpParent);

        if (NULL == parentObjId) {
            LOG_D("This is Ask key");
            status = loadAttestationKey(kpParent, false);
        } else {
            LOG_D("This is Auth key");
            status = tigerLoadKeyPair(parentObjId, kpParent);
        }
        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("Load key parent failed");

        TigerCounter_t counter = 0;
        status = tigerGetCounter(&counter);
        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("tigerGetCounter");

        unsigned char pemPubkeyData[RSA_PUB_DER_MAX_BYTES] = {0};
        ret = mbedtls_pk_write_pubkey_pem(&kp->ctx, pemPubkeyData, sizeof(pemPubkeyData));
        const size_t pemPubkeyDataLen = strlen((const char*) pemPubkeyData);
        TIGER_CHECK_MBEDTLS_RET_ZERO_BREAK("mbedtls_pk_write_pubkey_pem");

        // at this point some deallocation can be done
        tigerFreeKeyPair(kp);
        kp = NULL;

        //replace new line characters in pem
        const size_t nbNewLines = getNewLineCharactersCount(pemPubkeyData, pemPubkeyDataLen);
        const size_t pemNewLinesReplacedLen = pemPubkeyDataLen + nbNewLines;

        pemNewLinesReplaced = tzwMalloc(pemNewLinesReplacedLen + 1); // for '\0' character
        TIGER_CHECK_BUFFER_ALLOCATED_BREAK(pemNewLinesReplaced);

        ret = replaceNewLineCharacters(pemNewLinesReplaced, pemNewLinesReplacedLen, pemPubkeyData, pemPubkeyDataLen);
        if (0 != ret) {
            LOG_E("replaceNewLineCharacters failure.");
            status = TEE_ERROR_SHORT_BUFFER;
            break;
        }

        uint8_t cpuId[DEVICE_ID_SIZE] = {0};
        status = loadDeviceId(cpuId, DEVICE_ID_SIZE);
        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("loadDeviceId");
        logByteArrayString(cpuId, DEVICE_ID_SIZE, "DeviceId");

        // set json certificate data
        jsonCreate(&jsonCertObj);
        writeJsonCertificate(jsonCertObj, pemNewLinesReplaced, pemNewLinesReplacedLen, cpuId, DEVICE_ID_SIZE, counter, uid);
        uint32_t jsonCertDataLen = jsonGetLength(jsonCertObj);
        const char* jsonCertData = jsonGetPlain(jsonCertObj);

        logByteArrayString((const uint8_t*) jsonCertData, jsonCertDataLen, "JSON CertData");

        unsigned char hash[MBEDTLS_HASH_SIZE] = {0};
        unsigned char sig[TIGER_CERT_SIGN_LENGTH] = {0};

        // make hash
        ret = mbedtls_md(mbedtls_md_info_from_type(MBEDTLS_MD_SHA256), (const uint8_t*) jsonCertData, jsonCertDataLen, hash);
        TIGER_CHECK_MBEDTLS_RET_ZERO_BREAK("mbedtls_md");

        mbedtls_rsa_context* issuer_rsa_ctx = mbedtls_pk_rsa(kpParent->ctx);
        mbedtls_rsa_set_padding(issuer_rsa_ctx, MBEDTLS_RSA_PKCS_V21, MBEDTLS_MD_SHA256);
        ret = tigerMbedtlsRsaRsassaPssSign(issuer_rsa_ctx, tigerMbedtlsRng, NULL, MBEDTLS_RSA_PRIVATE,
                                 issuer_rsa_ctx->hash_id, 0, hash,
                                 TIGER_CERT_SIGN_SALT_LENGTH,
                                 sig);
        TIGER_CHECK_MBEDTLS_RET_ZERO_BREAK("tigerMbedtlsRsaRsassaPssSign");
        const size_t jsonCertDataSigLen = issuer_rsa_ctx->len;
        LOG_D("Signature length: %d", jsonCertDataSigLen);
        logByteArrayBase64(sig, jsonCertDataSigLen, "SIGNATURE");

        const size_t jsonFinalCertDataLength = sizeof(jsonCertDataLen) + jsonCertDataLen + jsonCertDataSigLen;
        TIGER_CHECK_CONDITION_RETURN(jsonFinalCertDataLength < UINT_MAX);
        if (*size < jsonFinalCertDataLength) {
            LOG_E("Final length certificate data exceeds output buffer length.");
            status = TEE_ERROR_BAD_PARAMETERS;
            break;
        }

        // setup output data
        tzwMemMove(p, &jsonCertDataLen, sizeof(jsonCertDataLen));
        p += sizeof(jsonCertDataLen);
        tzwMemMove(p, jsonCertData, jsonCertDataLen);
        p += jsonCertDataLen;
        tzwMemMove(p, sig, jsonCertDataSigLen);
        *size = (uint32_t) jsonFinalCertDataLength;
    } while(0);

    jsonRelease(jsonCertObj);
    tzwFree(pemNewLinesReplaced);
    tigerFreeKeyPair(kp);
    tigerFreeKeyPair(kpParent);
    LOG_FUNC_END;
    return status;
}

TigerKeyPair_t* tigerAllocateKeyPair(void) {
    TigerKeyPair_t* pkp = tzwMalloc(sizeof(TigerKeyPair_t));
    if (NULL != pkp) {
        mbedtls_pk_init(&pkp->ctx);
    }

    return pkp;
}

void tigerFreeKeyPair(TigerKeyPair_t* kp) {
    if (NULL != kp) {
        mbedtls_pk_free(&kp->ctx);
    }
    tzwFree(kp);
}

TEE_Result tigerKeyPairExists(const TigerObjectId_t* const objId) {

    TIGER_ASSERT(NULL != objId);
    TIGER_ASSERT(0 != tigerGetObjectIdLength(objId));

    TzwSfsObject_t handle = TEE_HANDLE_NULL;
    TEE_Result status = tzwOpenSfsObject((void*) tigerGetObjectIdData(objId), tigerGetObjectIdLength(objId),
                                         TEE_DATA_FLAG_ACCESS_READ,
                                         &handle);
    if (TEE_SUCCESS == status) {
        tzwCloseSfsObject(handle);
    }
    return status;
}

// ------------------ Internal functions implementation ----------------------------------------------------------------

int serializeKeyPair(const TigerKeyPair_t* const keypair, uint8_t* buf, size_t size) {
    LOG_FUNC_BEGIN;
    int ret;
    size_t len = 0;
    // serialize key pair into buffer
    MBEDTLS_ASN1_CHK_ADD(len, mbedtls_pk_write_key_der((mbedtls_pk_context* ) &keypair->ctx, buf, size));
    LOG_D("Size of der-encoded keypair is %d bytes", len);

    // write len and tag to denote the keypair data
    unsigned char* p = buf + size - len;
    unsigned char* start = buf;
    MBEDTLS_ASN1_CHK_ADD(len, mbedtls_asn1_write_len(&p, start, len));
    MBEDTLS_ASN1_CHK_ADD(len, mbedtls_asn1_write_tag(&p, start, MBEDTLS_ASN1_OCTET_STRING));
    LOG_FUNC_END;
    return ((int) len);
}

TEE_Result seekToKeyPair(TzwSfsObject_t handle, size_t* len) {
    return seekTo(handle, TKP_KEY_PAIR, len);
}

TEE_Result seekTo(TzwSfsObject_t handle, TigerKeyPart_t part, size_t* len) {

    LOG_FUNC_BEGIN;

    TIGER_CHECK_FUNCTION_ARGUMENT_RETURN(NULL != len);

    TEE_Result status = tzwSeekIntoSfsObject(handle, 0, TEE_DATA_SEEK_SET);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("tzwSeekIntoSfsObject");

    uint8_t buf[TIGER_TAG_AND_LEN_BUFFER_SIZE];
    *len = 0;
    for (int i = 0; i <= (int) part; ++i) {
        LOG_D("Iteration %d", i);
        status = tzwSeekIntoSfsObject(handle, *len, TEE_DATA_SEEK_CUR);
        if (TEE_SUCCESS != status) {
            logTeeError(status, __func__);
            break;
        }

        uint32_t count = 0;
        tzwReadFromSfsObject(handle, buf, TIGER_TAG_AND_LEN_BUFFER_SIZE, &count);
        if (count <= 0) {
            status = TEE_ERROR_NO_DATA;
            break;
        }

        unsigned char *p = buf;
        int ret = mbedtls_asn1_get_tag(&p, buf + count, len, MBEDTLS_ASN1_OCTET_STRING);
        if ((0 != ret) && (MBEDTLS_ERR_ASN1_OUT_OF_DATA != ret)) {
            logMbedtlsError(ret, __func__);
            status = TEE_ERROR_GENERIC;
            break;
        }

        // for TKP_KEY_PAIR and TKP_KEY_INFO we seek to position just after corresponding Tag-Length byte
        int32_t offset = (p - buf - TIGER_TAG_AND_LEN_BUFFER_SIZE);
        // for TKP_KEY_CHLD we seek to position just before Tag-Length byte, because child list is not wrapped by common TLV
        if (TKP_KEY_CHLD == i) {
            offset = -TIGER_TAG_AND_LEN_BUFFER_SIZE;
        }

        status = tzwSeekIntoSfsObject(handle, offset, TEE_DATA_SEEK_CUR);
        if (TEE_SUCCESS != status) {
            logTeeError(status, __func__);
            break;
        }
    };

    LOG_FUNC_END;

    return status;
}

TEE_Result writeJsonCertificate(JsonObject_t jsonCertObj, const uint8_t* pemNewLinesReplaced,
                                const uint32_t pemNewLinesReplacedSize,
                                const uint8_t* cpuId, const uint32_t cpuIdSize,
                                const uint64_t counter, uint32_t uid) {
    TEE_Result status = jsonOpenBrace(jsonCertObj);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteRawString");
    status = jsonWriteKeyBuffer(jsonCertObj, (const char*) "pub_key", pemNewLinesReplaced, pemNewLinesReplacedSize);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteKeyBuffer");
    status = jsonWriteComma(jsonCertObj);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteComma");
    status = jsonWriteKeyBuffer(jsonCertObj, (const char*) "cpu_id", cpuId, cpuIdSize);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteKeyBuffer");
    status = jsonWriteComma(jsonCertObj);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteComma");
    status = jsonWriteKeyUint64(jsonCertObj, (const char*) "counter", counter);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteKeyUint64");
    status = jsonWriteComma(jsonCertObj);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteComma");
    status = jsonWriteKeyUint32AsStr(jsonCertObj, (const char*) "uid", uid);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteKeyUint32AsStr");
    status = jsonCloseBrace(jsonCertObj);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteRawString");

    return status;
}