/*
 * 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: Mar 17, 2016
 * Author: Oleksii Kachkan <o.kachkan@samsung.com>
 * Brief: Attk key generation.
 */

#include <string.h>
#include <errno.h>

#include "TigerAttk.h"

#include <mbedtls/pk.h>

#include "TigerLogging.h"
#include "TzwRpmb.h"
#include "TzwHash.h"
#include "TzwMemory.h"
#include "TzwString.h"
#include "TigerMacros.h"
#include "TigerCore.h"
#include "TigerStorageUtils.h"
#include "TigerSskds.h"
#if !defined(TA_RELEASE) && defined(DEV_TA_DRK_DEBUG)
#include "qsee_fs.h"
#endif

#if defined(TIGER_BUILD_TESTS)
    #include "TigerTests.h"
#endif


#define TIGER_SKM_ALIAS                 "skm"
#define TIGER_TA_ALIAS_MAX_LEN          (32)

// Next few magic defines provided by DRK-team
/**
 * @brief Function to get 16 bits long Big-endian TLV-struct length stored in wrapped file
 */
#define GET_UINT16T_LENGTH(N, ARR) (uint16_t) ( ((uint16_t)ARR[N]) | ((uint16_t)ARR[N + 1] << 8) )

#define TAG_SIZE_BYTES (1)
#define LEN_SIZE_BYTES (2)
#define TL_SIZE_BYTES (TAG_SIZE_BYTES + LEN_SIZE_BYTES)

// Tags, used by DRK trustlet to store keys in wrapped file.
#define RSA_CERT_TAG      (uint8_t) 0x01
#define IV_TAG            (uint8_t) 0x02
#define KEY_TAG           (uint8_t) 0x03
#define TL_NAME_TAG       (uint8_t) 0x04
#define ATTRS_TAG         (uint8_t) 0x05


/**
 * @brief Parse attestation key's data, obtained from DeviceRootKeyManager.
 *
 * @param[in] drkData       - decrypted by DRK trustlet data.
 * @param[out] drkCert      - plain DRK-key certificate in DER-format.
 * @param[out] attkCert     - plain ATTK-key certificate in DER-format.
 * @param[out] attkPrivKey  - plain ATTK private key in DER-format.
 *
 * @return status of the operation, e.g. TEE_SUCCESS on success.
 */
static TEE_Result parseDrkData(IN TigerDynamicBuf_t* drkData, OUT TigerDynamicBuf_t* drkCert,
                        OUT TigerDynamicBuf_t* childCert, OUT TigerDynamicBuf_t* childKeyPair);

static TzwErrorCode_t getAttkHash(const uint8_t *bytes, const uint32_t dataLen, uint8_t* hash);

static TzwErrorCode_t checkAttk(uint8_t* keyData, uint32_t keyLen, uint8_t* expectedHash);

TEE_Result saveAttestationKey(const TciDrkDataMessage_t* wrappedDrkData) {
    LOG_FUNC_BEGIN;
    TIGER_ASSERT(wrappedDrkData != NULL)

    LOG_D("Clearing ATTK from RPMB.");
    rpmbClear();

    LOG_D("Clearing tigerfp TZ storage.");
    clearAll(true);

    TigerDynamicBuf_t* drkPlain = tigerAllocateDynamicBuf(wrappedDrkData->size);
    TIGER_CHECK_BUFFER_ALLOCATED_RETURN(drkPlain);

    TEE_Result status = TEE_SUCCESS;
    do {
        TigerDynamicBuf_t drkCertDer;
        TigerDynamicBuf_t attkCertDer;
        TigerDynamicBuf_t attkKeyPair;
        status = tzwUnwrapData(TIGER_SKM_ALIAS,
                          wrappedDrkData->bytes, wrappedDrkData->size,
                          drkPlain->bytes, &drkPlain->size);
        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("unwrapDrkData()");

#if !defined(TA_RELEASE) && defined(DEV_TA_DRK_DEBUG)
        LOG_D("++++++++++ DEV_TA_DRK_DEBUG defined ++++++++");
        int fp = 0;
        const char *fileName = NULL;
        uint32_t sz = 0;
#endif

        status = parseDrkData(drkPlain, &drkCertDer, &attkCertDer, &attkKeyPair);
        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("parseDrkData()");

#ifdef __DEV_DEBUG__
        logRawByteArrayHex(drkCertDer.bytes, drkCertDer.size, "DRK/Device CERT DER");
        logRawByteArrayHex(attkKeyPair.bytes, attkKeyPair.size, "ATTK/Private-key DER");
        logRawByteArrayHex(attkCertDer.bytes, attkCertDer.size, "ATTK/Service CERT DER");
#endif

        // save attk to RPMB
        status = saveAttk(attkKeyPair.bytes, attkKeyPair.size);
        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("saveAttk()");

#if !defined(TA_RELEASE) && defined(DEV_TA_DRK_DEBUG)
        fileName = "/efs/tigerfp/attk.priv";
        fp = open(fileName, O_RDWR|O_CREAT);
        if(fp < 0){
            LOG_E("open file failed %s, strerror: %s", fileName, strerror(errno));
            LOG_E("Missing track file: %s", fileName);
        }else{
            sz = write(fp, attkKeyPair.bytes, attkKeyPair.size);
            LOG_Raw("+++++++++++++write size %zd to %s++++++++++++++++++", sz, fileName);
            close(fp);
        }
#endif
        // Save ATTK Certificate to Secure storage
        status = saveToPersistentObject(TIGER_ATTK_CERT_ALIAS, attkCertDer.bytes, attkCertDer.size);
        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("saveCertificate(ATTK)");

#if !defined(TA_RELEASE) && defined(DEV_TA_DRK_DEBUG)
        fileName = "/efs/tigerfp/attk.cert";
        fp = open(fileName, O_RDWR|O_CREAT);
        if(fp < 0){
            LOG_E("open file failed %s, strerror: %s", fileName, strerror(errno));
            LOG_E("Missing track file: %s", fileName);
        }else{
            sz = write(fp, attkCertDer.bytes, attkCertDer.size);
            LOG_Raw("+++++++++++++write size %zd to %s++++++++++++++++++", sz, fileName);
            close(fp);
        }
#endif
        // Save DRK Certificate to secure storage
        status = saveToPersistentObject(TIGER_DRK_CERT_ALIAS, drkCertDer.bytes, drkCertDer.size);
        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("saveCertificate(DRK)");

#if !defined(TA_RELEASE) && defined(DEV_TA_DRK_DEBUG)
        fileName = "/efs/tigerfp/device.cert";
        fp = open(fileName, O_RDWR|O_CREAT);
        if(fp < 0){
            LOG_E("open file failed %s, strerror: %s", fileName, strerror(errno));
            LOG_E("Missing track file: %s", fileName);
        }else{
            sz = write(fp, drkCertDer.bytes, drkCertDer.size);
            LOG_Raw("+++++++++++++write size %zd to %s++++++++++++++++++", sz, fileName);
            close(fp);
        }
#endif

    } while (false);

    tigerFreeDynamicBuf(drkPlain);

    LOG_FUNC_END;
    return status;
}


TEE_Result loadAttestationKey(TigerKeyPair_t* keyPair, bool useNextAttk)  {
    LOG_FUNC_BEGIN;

    /// REMARK:
    /// From SOTER Adaptation for HAL and key master TA.pdf v2.4 page 13
    /// if is the current ATTK's digest equal to the saved digest? we use this copy
    /// else mark next
    S_VAR_NOT_USED(useNextAttk);

#if defined(TIGER_BUILD_TESTS)
        return getTestAttk(keyPair);
#endif

    TEE_Result status = TEE_SUCCESS;
    uint8_t attkRaw[MAX_EXPORT_SIZE_BYTES] = {0};

    do {
        status = rpmbRead(attkRaw, MAX_EXPORT_SIZE_BYTES);
        TIGER_CHECK_TEE_STATUS_SUCCESS_BREAK("rpmbRead");

#ifdef __DEV_DEBUG__
        logRawByteArrayHex(attkRaw, MAX_EXPORT_SIZE_BYTES, "ATTK RAW");
#endif

        // At front of buffer we have attk hash which always 32 bytes (TEE_ALG_SHA256_DIGEST_LENGTH)
        uint8_t* pSavedHash = attkRaw;
        // After hash we have key data length which always 4 bytes (sizeof(uint32_t))
        uint8_t* pKeyDataLength = attkRaw + TEE_ALG_SHA256_DIGEST_LENGTH;
        // After data len we have attk data in DER format
        uint8_t* pKeyData = attkRaw + TEE_ALG_SHA256_DIGEST_LENGTH + sizeof(uint32_t);

        const uint32_t keySize =  *((uint32_t*)pKeyDataLength);
        LOG_D("Read key data Len %u", keySize);

        if (keySize > MAX_EXPORT_SIZE_BYTES) {
            LOG_E("Bad Attk Size: %u", keySize);
            status = TEE_ERROR_CORRUPT_OBJECT;
            break;
        }

        // Verify that attk valid
        status = checkAttk(pKeyData, keySize, pSavedHash);
        if (TEE_SUCCESS != status) {
            logTeeError(status, "checkAttk()");
            continue;
        }

#ifdef __DEV_DEBUG__
        logRawByteArrayHex(pKeyData, keySize, "ATTK Key Data");
#endif

        LOG_D("Parse key-pair");
        int ret = mbedtls_pk_parse_key(tigerGetKeyPairContext(keyPair), pKeyData, keySize, NULL, 0);
        if (0 != ret) {
            logMbedtlsError(ret, "Failed to parse Attk");
            status = TEE_ERROR_CORRUPT_OBJECT;
            continue;
        }

        logRsaKeyPair(keyPair, "ATTK");

    } while(0);

    LOG_FUNC_END;
    return  status;
}


static TEE_Result parseDrkData(TigerDynamicBuf_t* drkData,
                        TigerDynamicBuf_t* drkCert,
                        TigerDynamicBuf_t* childCert,
                        TigerDynamicBuf_t* childKeyPair) {
    LOG_FUNC_BEGIN;
        /* In unwrapped file we have TLV-struct, which consist of 3 parts:
     *      - DRK certificate in DER format;
     *      - ATTK certificate in DER format;
     *      - ATTK private key in DER format
    */

    uint32_t index = 0;
    uint16_t keyLength = 0;

    // First is DRK certificate
    if (drkData->bytes[index] != RSA_CERT_TAG) {
        LOG_E("Incorrect blob (#1)");
        return TEE_ERROR_BAD_FORMAT;
    }

    keyLength = GET_UINT16T_LENGTH(index + TAG_SIZE_BYTES, drkData->bytes);
    drkCert->bytes = drkData->bytes + index + TL_SIZE_BYTES;
    drkCert->size = keyLength;
    index += (keyLength + TL_SIZE_BYTES);
    keyLength = 0;

    // Second is Attk certificate
    if (drkData->bytes[index] != RSA_CERT_TAG) {
        LOG_E("Incorrect blob (#2)");
        return TEE_ERROR_BAD_FORMAT;
    }

    keyLength = GET_UINT16T_LENGTH(index + TAG_SIZE_BYTES, drkData->bytes);
    childCert->bytes = drkData->bytes + index + TL_SIZE_BYTES;
    childCert->size = keyLength;

    index += (keyLength + TL_SIZE_BYTES);
    keyLength = 0;

    // Third is Attk private key
    if (drkData->bytes[index] != KEY_TAG) {
        LOG_E("Incorrect blob (#3)");
        return TEE_ERROR_BAD_FORMAT;
    }
    keyLength = GET_UINT16T_LENGTH(index + TAG_SIZE_BYTES, drkData->bytes);
    childKeyPair->bytes = drkData->bytes + index + TL_SIZE_BYTES;
    childKeyPair->size = keyLength;

    LOG_D("All keys was parsed!");
    LOG_FUNC_END;

    return TEE_SUCCESS;
}

TEE_Result saveAttk(const uint8_t* keyData, const uint32_t keyLength) {
    LOG_FUNC_BEGIN;

    // Size: < sha256 hash(32 bytes)> + < length of ATTK_pri (4 bytes)> +  <ATTK_pri>
    size_t packedDataSizeTemp = TEE_ALG_SHA256_DIGEST_LENGTH + sizeof(uint32_t) + keyLength;
    TIGER_CHECK_CONDITION_RETURN(packedDataSizeTemp < UINT_MAX);

    const uint32_t packedDataSize = (uint32_t) packedDataSizeTemp;

    // prepare Attk and it's hash
    uint8_t* attkPackedData = tzwMalloc(packedDataSize);
    TIGER_CHECK_BUFFER_ALLOCATED_RETURN(attkPackedData);

    uint8_t attkSha256Hash [TEE_ALG_SHA256_DIGEST_LENGTH] = {0};
    TEE_Result status = getAttkHash(keyData, keyLength, attkSha256Hash);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("Calculate ATTK sha256");

    // Partition: < sha256 hash(32 bytes)> + < length of ATTK_pri (4 bytes)> +  <ATTK_pri>
    tzwMemMove(attkPackedData, attkSha256Hash, sizeof(attkSha256Hash));
    tzwMemMove(attkPackedData + sizeof(attkSha256Hash), (void*)&keyLength, sizeof(keyLength));
    tzwMemMove(attkPackedData + sizeof(attkSha256Hash) + sizeof(keyLength), keyData, keyLength);

    status = rpmbWrite(attkPackedData, packedDataSize);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("rpmbWrite()");

    tzwFree(attkPackedData);
    LOG_FUNC_END;
    return  status;
}

static TzwErrorCode_t getAttkHash(const uint8_t* bytes, const uint32_t dataLen, uint8_t* hash) {
    LOG_FUNC_BEGIN;
    LOG_D("dataLen: %u", dataLen);

    TzwHashOperation_t digestOperation = TEE_HANDLE_NULL;
    size_t hashLen = TEE_ALG_SHA256_DIGEST_LENGTH;

    TzwErrorCode_t status = tzwAllocateHashOperation(&digestOperation, TEE_ALG_SHA256);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("tzwAllocateHashOperation()");

    status = tzwHashFinalize(digestOperation, (void*) bytes, dataLen, hash, &hashLen);
    if (TEE_SUCCESS != status) {
        logTeeError(status, "tzwHashFinalize()");
    }

    tzwFreeHashOperation(digestOperation);

    LOG_FUNC_END;
    return status;
}

static TzwErrorCode_t checkAttk(uint8_t* keyData, uint32_t keyLen, uint8_t* expectedHash) {
    LOG_FUNC_BEGIN;

    // For calculate  sha256(ATTK_pri)
    uint8_t actualHash [TEE_ALG_SHA256_DIGEST_LENGTH] = {0};
    TzwErrorCode_t status = getAttkHash(keyData, keyLen, actualHash);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("getAttkHash");

    logByteArrayHex(expectedHash, TEE_ALG_SHA256_DIGEST_LENGTH, "Expected hash");
    logByteArrayHex(actualHash, TEE_ALG_SHA256_DIGEST_LENGTH, "Actual hash");


    // Verify that attk valid
    int32_t ret = tzwMemCompare((void*) expectedHash, (void*) actualHash, TEE_ALG_SHA256_DIGEST_LENGTH);
    if (0 != ret) {
        LOG_E("Failed to verify ATTK");
        return TEE_ERROR_CORRUPT_OBJECT;
    }
    LOG_D("Attk hash is correct! ");


    LOG_FUNC_END;
    return TEE_SUCCESS;
}