/*
 * app_cipher.c (copied from TZ_Vendor_tl.c (tima_common) to remove the code dependency in TIMA path)
 */

#include <stdio.h>
#include <tee_internal_api.h>

#include "app_cipher.h"
#include "app_core.h"

#include "icccOperations_v4.h"

uint32_t Iccc_digest_SHA256(
    uint8_t *messageData,
    uint32_t messageLen,
    uint8_t *digest,
    uint32_t *pDigestLen
)
{
    uint32_t ret = ICCC_SUCCESS;
    TEE_OperationHandle opHandle = NULL;
    TEE_Result teeRet = TZ_API_OK;

    ICCC_LOG("generate digest");

    if (messageData == NULL || messageLen == 0) {
        ICCC_LOG("invalid input data buffer");
        return ICCC_ERROR_ATTESTATION_FAILED;
    }

    if (digest == NULL || pDigestLen == NULL) {
        ICCC_LOG("invalid output digest data buffer");
        return ICCC_ERROR_ATTESTATION_FAILED;
    }

    if (*pDigestLen < SHA256_DIGEST_LENGTH) {
        ICCC_LOG("pDigestLen < SHA256_DIGEST_LENGTH");
        return ICCC_ERROR_ATTESTATION_FAILED;
    }

    teeRet = TEE_AllocateOperation(&opHandle, TEE_ALG_SHA256, TEE_MODE_DIGEST, TEE_MAX_KEY_SIZE);
    if (TEE_SUCCESS != ret) {
        ICCC_LOG("Error: TEE_AllocateOperation failed with ret = %x", teeRet);
        ret = ICCC_ERROR_ATTESTATION_FAILED;
        goto exit;
    }

    teeRet = TEE_DigestDoFinal(opHandle, messageData, messageLen, digest, pDigestLen);
    if (TEE_SUCCESS != ret) {
        ICCC_LOG("Error: TEE_DigestDoFinal() failed with ret = %x", teeRet);
        ret = ICCC_ERROR_ATTESTATION_FAILED;
        goto exit;
    }

    ICCC_LOG("TEE_DigestDoFinal() success digest len = %d", *pDigestLen);
    // TEE_DigestDoFinal() has a bug with which pDigestLen is not returned correctly.
    // So here pDigestLen is set mannually
    *pDigestLen = SHA256_DIGEST_LENGTH;

    ret = ICCC_SUCCESS;

exit:
    if (NULL!=opHandle) {
        TEE_FreeOperation(opHandle);
    }

    return ret;
}

uint32_t Iccc_sign_CKM_SHA256_RSA_PKCS_PSS(
    uint8_t *keyMod,
    uint32_t keyModLen,
    uint8_t *keyPubExp,
    uint32_t keyPubExpLen,
    uint8_t *keyPriExp,
    uint32_t keyPriExpLen,
    uint8_t *messageData,
    uint32_t messageLen,
    uint8_t *signature,
    uint32_t *pSigLen
)
{
    TEE_OperationHandle opHandle = NULL;
    TEE_ObjectHandle keyHandle = NULL;
    uint32_t ret = TZ_API_OK;
    uint8_t digest[SHA256_DIGEST_LENGTH] = {0};
    uint32_t digestLen = sizeof(digest);

    TEE_Attribute attrs[SIGN_ATTR_COUNT];

    ICCC_LOG("sign using RSA key");

    if (keyMod == NULL || keyModLen == 0) {
        ICCC_LOG("invalid input key modulus");
        return TZ_API_ERROR;
    }

    //ICCC_LOG("keyModLen %d", keyModLen);
    if (keyModLen > MAX_SIGNATURE_SIZE) {
        ICCC_LOG("signature length exceeds MAX_SIGNATURE_SIZE");
        return TZ_API_ERROR;
    }

    if (keyPubExp == NULL || keyPubExpLen == 0) {
        ICCC_LOG("invalid input key public exponent");
        return TZ_API_ERROR;
    }

    if (keyPriExp == NULL || keyPriExpLen == 0) {
        ICCC_LOG("invalid input key private exponent");
        return TZ_API_ERROR;
    }

    if (messageData == NULL || messageLen == 0) {
        ICCC_LOG("invalid input signing data");
        return TZ_API_ERROR;
    }

    if (signature == NULL || pSigLen == NULL|| *pSigLen == 0) {
        ICCC_LOG("invalid output signature data");
        return TZ_API_ERROR;
    }

    ret = Iccc_digest_SHA256(messageData, messageLen, digest, &digestLen);
    if (TZ_API_OK != ret) {
        ICCC_LOG("Iccc_digest_SHA256() returns an error code 0x%x", ret);
        return TZ_API_ERROR;
    }

    ret = TEE_AllocateTransientObject(TEE_TYPE_RSA_KEYPAIR, 8*keyModLen, &keyHandle);
    if (TEE_SUCCESS != ret) {
        ICCC_LOG("TEE_AllocateTransientObject failed with ret = %d", ret);
        goto exit;
    }
    TEE_InitRefAttribute(&attrs[0], TEE_ATTR_RSA_MODULUS, keyMod, keyModLen);
    TEE_InitRefAttribute(&attrs[1], TEE_ATTR_RSA_PUBLIC_EXPONENT, keyPubExp, keyPubExpLen);
    TEE_InitRefAttribute(&attrs[2], TEE_ATTR_RSA_PRIVATE_EXPONENT, keyPriExp, keyPriExpLen);

    ret = TEE_PopulateTransientObject(keyHandle, attrs, SIGN_ATTR_COUNT);
    if (TEE_SUCCESS != ret) {
        ICCC_LOG("TEE_PopulateTransientObject failed with ret = %d", ret);
        goto exit;
    }

    ret = TEE_AllocateOperation(&opHandle, TEE_ALG_RSASSA_PKCS1_PSS_MGF1_SHA256, TEE_MODE_SIGN, 8*keyModLen);
    if (TEE_SUCCESS != ret) {
        ICCC_LOG("TEE_AllocateOperation failed with ret = %d", ret);
        goto exit;
    }

    ret = TEE_SetOperationKey(opHandle, keyHandle);
    if (TEE_SUCCESS != ret) {
        ICCC_LOG("TEE_SetOperationKey failed with ret = %d", ret);
        goto exit;
    }

    ret = TEE_AsymmetricSignDigest(opHandle, NULL, 0, digest, digestLen, signature, pSigLen);
    if (TEE_SUCCESS != ret) {
        ICCC_LOG("TEE_AsymmetricSignDigest failed with ret = %d", ret);
        goto exit;
    }

    ICCC_LOG("TEE_AsymmetricSignDigest() success");
    ret = TZ_API_OK;

exit:

    if (keyHandle != NULL) {
        TEE_FreeTransientObject(keyHandle);
    }
    if (opHandle != NULL) {
        TEE_FreeOperation(opHandle);
    }
    return ret;
}