/*
 * 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)
 *
 * Created on: May 19, 2016
 * Author: Konstyantyn Volobuyev <k.volobuyev@samsung.com>
 * Brief: Signature operation.
 */

#include "TigerSignMessage.h"

#include <string.h>
#include <stdio.h>
#include <inttypes.h>
#include <mbedtls/pk.h>
#include "mbedtls/asn1.h"
#include "mbedtls/asn1write.h"

#include "TigerCounterDataStore.h"
#include "TigerKeyDataStore.h"
#include "TigerLogging.h"
#include "TigerCore.h"
#include "TigerAttk.h"
#include "TigerMacroses.h"
#include "TigerMbedTlsExt.h"
#include "TigerSskds.h"
#include "TigerStorageUtils.h"
#include "TigerJson.h"
#include "TigerSession.h"
#include "TigerFingerprintIdTable.h"
#include "TzwMemory.h"
#include "TzwString.h"

#define MAX_DIGEST_SIZE     (32)
#define TIGER_MSG_SIGN_SALT_LENGTH  20 /*bytes*/

#define MBEDTLS_HASH_SIZE 64
#define TIGER_FINISH_SIGN_LENGTH      256

extern const char* TEE_NAME;
extern const char* TEE_VERSION;
extern const char* FINGERPRINT_SENSOR_NAME;
extern const char* FINGERPRINT_SENSOR_VERSION;

const char* FORMAT_BEGIN_KEY = "{\"raw\":\"";

static TEE_Result tigerWriteJsonSignature(const uint32_t fid, TigerSessionContext* context , uint8_t* buf, uint32_t* buflen);
static TEE_Result tigerGetLatestAuthResult(uint32_t *soterFid);


TEE_Result initSign(const TciInitSignMessage_t* const request, TciInitSignResponse_t* response) {
    LOG_FUNC_BEGIN;

    TIGER_CHECK_FUNCTION_ARGUMENT_RETURN(request != NULL);
    TIGER_CHECK_FUNCTION_ARGUMENT_RETURN(response != NULL);

    TigerSessionContext* context = tigerAllocateSessionContext();

    TEE_Result status = tigerLoadSessionContext(context);

    if (TEE_ERROR_ITEM_NOT_FOUND == status || TEE_ERROR_CORRUPT_OBJECT == status) {
        context->alias      = request->kc.alias;
        context->uid        = request->kc.uid;
        context->challenge  = request->challenge;
        context->sessionId  = tigerGenerateSessionId(request->kc.alias,request->challenge,
                                                     request->kc.uid);

        context->genTime    = tigerGetCurrTime();
        status = tigerSaveSessionContext(context);
        //return status;
    } else if(TEE_SUCCESS == status) {

        bool compareResult = tigerCompareSessionContext(context, request->kc.alias, request->challenge,
                                                      request->kc.uid); //0 same, 1 different
        bool isTimeout = tigerIsSessionTimout(context); //1 timeout
        LOG_I("initSign isTimeount(%d) compareResult(%d)", isTimeout, compareResult);

        if( compareResult || isTimeout) {
            tigerRemoveSessionObject();
            context->alias      = request->kc.alias;
            context->uid        = request->kc.uid;
            context->challenge  = request->challenge;
            context->sessionId  = tigerGenerateSessionId(request->kc.alias,request->challenge,
                                                         request->kc.uid);
            context->genTime    = tigerGetCurrTime();
            status = tigerSaveSessionContext(context);
        }

    }
    response->sessionId = context->sessionId;
    tigerFreeSessionContext(context);
    LOG_FUNC_END;
    return status;
}

TEE_Result finish(const TciSignFinishRequest_t* const pRequest, TciFinishSignData_t* signData) {
    LOG_FUNC_BEGIN;
    LOG_I("finish start session ID(%"PRIu64")", pRequest->sessionId);

    TIGER_CHECK_FUNCTION_ARGUMENT_RETURN(pRequest != NULL);
    TIGER_CHECK_FUNCTION_ARGUMENT_RETURN(signData != NULL);

    TigerSessionContext* context = tigerAllocateSessionContext();
    TEE_Result status = tigerLoadSessionContext(context);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("tigerLoadSessionContext");

    if(TEE_SUCCESS == status) {

        do {
            if(context->sessionId != pRequest->sessionId ) {//0 same, 1 different
                return TEE_ERROR_ITEM_NOT_FOUND;
            }
            if(tigerIsSessionTimout(context)) {//1 timeout
                LOG_I("tigerIsSessionTimout timeout is true");
                return status;
            }


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

            uint32_t soterFid;
            status = tigerGetLatestAuthResult(&soterFid);
            LOG_I("finish soterFid(%u)", soterFid);
            TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("tigerGetLatestAuthResult");

            signData->size = MAX_FINISH_SIZE_BYTES;
            // Attach specified data to end of the message and hash it
            status = tigerWriteJsonSignature(soterFid, context, signData->bytes,
                                        &signData->size);
            TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("tigerWriteJsonSignature");
        } while(0);

        LOG_I("finish json generate successfully");
    }

    tigerFreeSessionContext(context);
    tigerRemoveSessionObject();
    tigerFPTableSave();
    LOG_FUNC_END;
    return status;
}

TEE_Result tigerWriteJsonSignature(const uint32_t fid, TigerSessionContext* context, uint8_t* buf, uint32_t* buflen) {
    LOG_FUNC_BEGIN;

    TIGER_CHECK_FUNCTION_ARGUMENT_RETURN(buflen != NULL);

    TigerObjectId_t* objId = NULL;
    TigerKeyPair_t* kp = NULL;
    TEE_Result status = tigerCreateKeyPairObjectId(&(context->alias), context->uid, &objId);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("tigerCreateKeyPairObjectId");

    do {
        kp = tigerAllocateKeyPair();
        TIGER_CHECK_BUFFER_ALLOCATED_BREAK(kp);

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

        uint8_t deviceId[DEVICE_ID_SIZE] = {0};
        status = loadDeviceId(deviceId, DEVICE_ID_SIZE);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("loadDeviceId");

        uint64_t counter = 0;
        status = tigerGetIncCounter(&counter);
        if (TEE_SUCCESS != status) {
            counter = 0;
        }

        // adding of sign message parameters to json object
        JsonObject_t jsonParamsObj = NULL;
        jsonCreate(&jsonParamsObj);
        status = jsonWriteRawString(jsonParamsObj, FORMAT_BEGIN_KEY);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteRawString");
        status = jsonWriteRawString(jsonParamsObj, (char *)context->challenge.value);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteRawString");
        status = jsonWriteRawString(jsonParamsObj, "\"");
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteRawString");
        status = jsonWriteComma(jsonParamsObj);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteComma");
        status = jsonWriteKeyUint32AsStr(jsonParamsObj, "fid", fid);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteKeyUint32AsStr");
        status = jsonWriteComma(jsonParamsObj);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteRawString");
        status = jsonWriteKeyUint32(jsonParamsObj, "counter", counter);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteKeyUint32");
        status = jsonWriteComma(jsonParamsObj);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteComma");
        status = jsonWriteKeyString(jsonParamsObj, "tee_n", TEE_NAME);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteKeyString");
        status = jsonWriteComma(jsonParamsObj);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteComma");
        status = jsonWriteKeyString(jsonParamsObj, "tee_v", TEE_VERSION);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteKeyString");
        status = jsonWriteComma(jsonParamsObj);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteComma");
        status = jsonWriteKeyString(jsonParamsObj, "fp_n", FINGERPRINT_SENSOR_NAME);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteKeyString");
        status = jsonWriteComma(jsonParamsObj);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteComma");
        status = jsonWriteKeyString(jsonParamsObj, "fp_v", FINGERPRINT_SENSOR_VERSION);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteKeyString");
        status = jsonWriteComma(jsonParamsObj);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteComma");
        status = jsonWriteKeyBuffer(jsonParamsObj, "cpu_id", deviceId, DEVICE_ID_SIZE);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteKeyBuffer");
        status = jsonWriteComma(jsonParamsObj);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteRawString");
        status = jsonWriteKeyUint32AsStr(jsonParamsObj, "uid", context->uid);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteKeyUint32AsStr");
        status = jsonCloseBrace(jsonParamsObj);
        TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("jsonWriteRawString");

        uint32_t jsonSignDataLen = jsonGetLength(jsonParamsObj);
        const char* jsonSignData = jsonGetPlain(jsonParamsObj);

#ifdef __DEV_DEBUG__
        logByteArrayString((const uint8_t*) jsonSignData, jsonSignDataLen, "JSON CertData");
#endif

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

        int ret = mbedtls_md(mbedtls_md_info_from_type(MBEDTLS_MD_SHA256), (const uint8_t*) jsonSignData, jsonSignDataLen, hash);
        TIGER_CHECK_MBEDTLS_RET_ZERO_BREAK("mbedtls_md");

        mbedtls_rsa_context* rsa_ctx = mbedtls_pk_rsa(*((mbedtls_pk_context* )tigerGetKeyPairContext(kp)));
        mbedtls_rsa_set_padding(rsa_ctx, MBEDTLS_RSA_PKCS_V21, MBEDTLS_MD_SHA256);
        ret = tigerMbedtlsRsaRsassaPssSign(rsa_ctx, tigerMbedtlsRng, NULL, MBEDTLS_RSA_PRIVATE, rsa_ctx->hash_id,
                        0, hash, TIGER_MSG_SIGN_SALT_LENGTH, sig);
        TIGER_CHECK_MBEDTLS_RET_ZERO_BREAK("tigerMbedtlsRsaRsassaPssSign");
        const size_t jsonSignDataSigLen = rsa_ctx->len;
#ifdef __DEV_DEBUG__
        LOG_D("Signature length: %d", jsonSignDataSigLen);
        logByteArrayBase64(sig, jsonSignDataSigLen, "SIGNATURE");
#endif

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

        tzwMemMove(buf, &jsonSignDataLen, sizeof(jsonSignDataLen));
        buf += sizeof(jsonSignDataLen);
        tzwMemMove(buf, jsonSignData, jsonSignDataLen);
        buf += jsonSignDataLen;
        tzwMemMove(buf, sig, jsonSignDataSigLen);
        *buflen = (uint32_t) jsonFinalCertDataLength;

        jsonRelease(jsonParamsObj);
    } while(0);

    LOG_FUNC_END;
    return status;
}

TEE_Result tigerGetLatestAuthResult(uint32_t* soterFid) {
    LOG_FUNC_BEGIN;

    uint8_t realFid[MAX_BIO_EXT_DATA_SIZE] = {0};
    uint32_t realFidBufferSize = MAX_BIO_EXT_DATA_SIZE;
    uint32_t fpIndex = 0;

    // communicate with fp TA to check last authorized operation
    TEE_Result status = getLatestAuthResult(realFid, &realFidBufferSize, &fpIndex);
    if (TEE_SUCCESS != status) {
        LOG_I("getLatestAuthResult failed, status=%" PRIu32"", status);
        return TEE_ERROR_ACCESS_DENIED;
    }

    // setup soter fid into session context
    status = tigerFPTableGetUniqueId(realFid, realFidBufferSize, fpIndex, soterFid);
    if (TEE_SUCCESS != status) {
       LOG_E("tigerFPTableGetUniqueId failed, status=%" PRIu32"", status);
       return TEE_ERROR_ACCESS_DENIED;
    }
    LOG_FUNC_END;
    return TEE_SUCCESS;
}