#include "DDAR_TZ_Vendor_tl.h"
#include "SCrypto_Crypto.h"

#include "openssl/rsa.h"
#include "openssl/evp.h"
#include "openssl/obj_mac.h"
#include <openssl/ec.h>
#include <openssl/ecdh.h>
#include <openssl/ecdsa.h>
#include <openssl/aes.h>
#include <openssl/rand.h>
#include <openssl/bn.h>

#include "utils.h"

//#if defined(CONFIG_SMDK8895) || defined(CONFIG_SMDK9810) || defined(CONFIG_SMDK7870) || defined(CONFIG_SMDK7885)
#if 1
#define _EVP_CIPHER_CTX_init        EVP_CIPHER_CTX_init

#define _EVP_CipherInit             EVP_CipherInit
#define _EVP_Cipher                 EVP_Cipher

#define _EVP_EncryptInit_ex         EVP_EncryptInit_ex
#define _EVP_EncryptUpdate          EVP_EncryptUpdate
#define _EVP_EncryptFinal_ex        EVP_EncryptFinal_ex
#define _EVP_DecryptInit_ex         EVP_DecryptInit_ex
#define _EVP_DecryptUpdate          EVP_DecryptUpdate
#define _EVP_DecryptFinal_ex        EVP_DecryptFinal_ex

#define _EVP_CIPHER_CTX_ctrl        EVP_CIPHER_CTX_ctrl
#define _EVP_CIPHER_CTX_set_padding EVP_CIPHER_CTX_set_padding
#define _EVP_CIPHER_CTX_cleanup     EVP_CIPHER_CTX_cleanup

#else
#define _EVP_CIPHER_CTX_init        FIPS_cipher_ctx_init

#define _EVP_CipherInit             FIPS_CipherInit
#define _EVP_Cipher                 FIPS_Cipher

#define _EVP_EncryptInit_ex         FIPS_EncryptInit_ex
#define _EVP_EncryptUpdate          FIPS_EncryptUpdate
#define _EVP_EncryptFinal_ex        FIPS_EncryptFinal_ex
#define _EVP_DecryptInit_ex         FIPS_DecryptInit_ex
#define _EVP_DecryptUpdate          FIPS_DecryptUpdate
#define _EVP_DecryptFinal_ex        FIPS_DecryptFinal_ex

#define _EVP_CIPHER_CTX_ctrl        FIPS_CIPHER_CTX_ctrl
#define _EVP_CIPHER_CTX_set_padding FIPS_CIPHER_CTX_set_padding
#define _EVP_CIPHER_CTX_cleanup     FIPS_CIPHER_CTX_cleanup

#endif

uint32_t DDAR_TZ_pbkdf_hmac256(uint8_t* pwd, uint32_t pwd_len, uint8_t* salt, uint32_t salt_len,
                  uint32_t iterations, uint8_t* out, uint32_t out_len)
{
    uint32_t ret = DDAR_TZ_API_ERROR;

    if(!PKCS5_PBKDF2_HMAC((const char *)pwd, pwd_len, (const unsigned char *)salt, salt_len,
                                iterations, EVP_sha256(),
                                out_len, out)) {
        DUALDAR_LOG("%s: PKCS5_PBKDF2_HMAC(pwd,salt, out) failed", __FUNCTION__);
        goto exit;
    }

    ret = DDAR_TZ_API_OK;

exit:
    return ret;
}


uint32_t DDAR_TZ_aes_gcm256_encrypt(uint8_t *plaintext, uint32_t plaintext_len,
            uint8_t *aad, uint32_t aad_len, uint8_t *key, uint8_t *iv,
            uint8_t *ciphertext, uint32_t *ciphertextlen,
            uint8_t *tag)
{
    EVP_CIPHER_CTX ctx;
    int tmp_len;
    int tmp_ciphertext_len;
    uint32_t ret = DDAR_TZ_API_ERROR;

    /* initialise the context */
    EVP_CIPHER_CTX_init(&ctx);

    /* Initialise the encryption operation. */
    if(1 != EVP_EncryptInit_ex(&ctx, EVP_aes_256_gcm(), NULL, NULL, NULL)) {
        DUALDAR_LOG("%s: EVP_EncryptInit_ex(EVP_aes_256_gcm) failed", __FUNCTION__);
        goto exit;
    }

    /* Set IV length if default 12 bytes (96 bits) is not appropriate */
    if(1 != EVP_CIPHER_CTX_ctrl(&ctx, EVP_CTRL_GCM_SET_IVLEN, AES256_GCM_IV_SIZE, NULL)){
        DUALDAR_LOG("%s: EVP_CIPHER_CTX_ctrl(EVP_CTRL_GCM_SET_IVLEN) failed", __FUNCTION__);
        goto exit;
    }

    /* Initialise key and IV */
    if(1 != EVP_EncryptInit_ex(&ctx, NULL, NULL, key, iv)) {
        DUALDAR_LOG("%s: EVP_EncryptInit_ex(key,iv) failed", __FUNCTION__);
        goto exit;
    }

    EVP_CIPHER_CTX_set_padding(&ctx, 0);

    /* Provide any AAD data. This can be called zero or more times as
     * required
     */
    if(1 != EVP_EncryptUpdate(&ctx, NULL, &tmp_len, aad, aad_len)) {
        DUALDAR_LOG("%s: EVP_EncryptUpdate(aad) failed", __FUNCTION__);
        goto exit;
    }

    /* Provide the message to be encrypted, and obtain the encrypted output.
     * EVP_EncryptUpdate can be called multiple times if necessary
     */
    if(1 != EVP_EncryptUpdate(&ctx, ciphertext, &tmp_len, plaintext, plaintext_len)) {
        DUALDAR_LOG("%s: EVP_EncryptUpdate(ciphertext, plaintext) failed", __FUNCTION__);
        goto exit;
    }

    tmp_ciphertext_len = tmp_len;

    /* Finalise the encryption. Normally ciphertext bytes may be written at
     * this stage, but this does not occur in GCM mode
     */
    if(1 != EVP_EncryptFinal_ex(&ctx, ciphertext + tmp_len, &tmp_len)) {
        DUALDAR_LOG("%s: EVP_EncryptFinal_ex(ciphertext, plaintext) failed", __FUNCTION__);
        goto exit;
    }

    tmp_ciphertext_len += tmp_len;

    /* Get the tag */
    if(1 != EVP_CIPHER_CTX_ctrl(&ctx, EVP_CTRL_GCM_GET_TAG, AES256_GCM_TAG_SIZE, tag)) {
        DUALDAR_LOG("%s: EVP_CIPHER_CTX_ctrl(EVP_CTRL_GCM_GET_TAG, tag) failed", __FUNCTION__);
        goto exit;
    }

    *ciphertextlen = tmp_ciphertext_len;

    ret = DDAR_TZ_API_OK;

exit:
    _EVP_CIPHER_CTX_cleanup(&ctx);
    return ret;

}

uint32_t DDAR_TZ_aes_gcm256_decrypt(uint8_t *ciphertext, uint32_t ciphertext_len, uint8_t *aad,
            uint32_t aad_len, uint8_t *tag, uint8_t *key, uint8_t *iv,
            uint8_t *plaintext, uint32_t *plaintextlen)
{
    uint32_t ret = DDAR_TZ_API_ERROR;
    EVP_CIPHER_CTX ctx;
    int tmp_len;
    int tmp_plaintext_len;

    /* initialise the context */
    EVP_CIPHER_CTX_init(&ctx);

    EVP_CIPHER_CTX_set_padding(&ctx, 0);

    /* Initialise the decryption operation. */
    if(!EVP_DecryptInit_ex(&ctx, EVP_aes_256_gcm(), NULL, NULL, NULL)) {
        DUALDAR_LOG("%s: EVP_DecryptInit_ex(EVP_aes_256_gcm) failed", __FUNCTION__);
        goto exit;
    }

    /* Set IV length. Not necessary if this is 12 bytes (96 bits) */
    if(!EVP_CIPHER_CTX_ctrl(&ctx, EVP_CTRL_GCM_SET_IVLEN, AES256_GCM_IV_SIZE, NULL)) {
        DUALDAR_LOG("%s: EVP_CIPHER_CTX_ctrl(EVP_CTRL_GCM_SET_IVLEN) failed", __FUNCTION__);
        goto exit;
    }

    /* Initialise key and IV */
    if(!EVP_DecryptInit_ex(&ctx, NULL, NULL, key, iv)) {
        DUALDAR_LOG("%s: EVP_DecryptInit_ex(key, iv) failed", __FUNCTION__);
        goto exit;
    }

    /* Set expected tag value. Works in OpenSSL 1.0.1d and later */
    if(!EVP_CIPHER_CTX_ctrl(&ctx, EVP_CTRL_GCM_SET_TAG, AES256_GCM_TAG_SIZE, tag)) {
        DUALDAR_LOG("%s: EVP_CIPHER_CTX_ctrl(EVP_CTRL_GCM_SET_TAG, AES256_GCM_TAG_SIZE) failed", __FUNCTION__);
        goto exit;
    }

    /* Provide any AAD data. This can be called zero or more times as
     * required
     */
    if(!EVP_DecryptUpdate(&ctx, NULL, &tmp_len, aad, aad_len)) {
        DUALDAR_LOG("%s: EVP_DecryptUpdate(aad) failed", __FUNCTION__);
        goto exit;
    }

    /* Provide the message to be decrypted, and obtain the plaintext output.
     * EVP_DecryptUpdate can be called multiple times if necessary
     */
    if(!EVP_DecryptUpdate(&ctx, plaintext, &tmp_len, ciphertext, ciphertext_len)) {
        DUALDAR_LOG("%s: EVP_DecryptUpdate(plaintext, ciphertext) failed", __FUNCTION__);
        goto exit;
    }

    tmp_plaintext_len = tmp_len;

    /* Set expected tag value. Works in OpenSSL 1.0.1d and later */
    if(!EVP_CIPHER_CTX_ctrl(&ctx, EVP_CTRL_GCM_SET_TAG, AES256_GCM_TAG_SIZE, tag)) {
        DUALDAR_LOG("%s: EVP_CIPHER_CTX_ctrl(EVP_CTRL_GCM_SET_TAG, AES256_GCM_TAG_SIZE) failed", __FUNCTION__);
        goto exit;
    }

    /* Finalise the decryption. A positive return value indicates success,
     * anything else is a failure - the plaintext is not trustworthy.
     */
    if (!EVP_DecryptFinal_ex(&ctx, plaintext + tmp_len, &tmp_len)) {
        DUALDAR_LOG("%s: EVP_DecryptFinal_ex() failed", __FUNCTION__);
        goto exit;
    }

    /* Success */
    tmp_plaintext_len += tmp_len;
    *plaintextlen = tmp_plaintext_len;

    ret = DDAR_TZ_API_OK;

exit:
    _EVP_CIPHER_CTX_cleanup(&ctx);
    return ret;
}

//////////////////////////////////////
/////////////////////////////////////

uint32_t ecdh_prep(void **localkey, uint8_t *pubkey, uint32_t *pubkey_len) {
    EC_KEY *key;
    uint32_t ret = DDAR_TZ_API_ERROR;

    DUALDAR_LOG("ecdh_prep : start");

    /* Create an Elliptic Curve Key object and set it up to use the ANSI X9.62 Prime 256v1 curve */
    if(NULL == (key = EC_KEY_new_by_curve_name(NID_secp521r1))) {
       DUALDAR_LOG("EC_KEY_new_by_curve_name: failed");
       goto exit;
    }

    /* Generate the private and public key */
    if(1 != EC_KEY_generate_key(key)) {
        DUALDAR_LOG("EC_KEY_generate_key: failed");
        ecdh_keyfree(key);
        goto exit;
    }

    const EC_POINT *pub_key = EC_KEY_get0_public_key(key);

    *pubkey_len = EC_POINT_point2oct(EC_KEY_get0_group(key), pub_key, POINT_CONVERSION_UNCOMPRESSED,
                pubkey, *pubkey_len, NULL);

    *localkey = key;

    DUALDAR_LOG("ecdh_prep : end");

    ret = DDAR_TZ_API_OK;

exit:
    return ret;
}

void ecdh_keyfree(void *localkey) {
    DUALDAR_LOG("ecdh_keyfree");
    EC_KEY_free((EC_KEY *)localkey);
}

uint32_t ecdh_sessionkey_gen(void *localkey, uint8_t *pubkey, uint32_t pubkey_len,
        uint8_t *secret, uint32_t *secret_len)
{
    EC_KEY *key = localkey;
    EC_KEY    *peerkey = NULL;
    EC_POINT  *peerpub_key;
    const EC_GROUP *peergroup;

    int field_size;
    uint32_t ret = DDAR_TZ_API_ERROR;

    DUALDAR_LOG("ecdh_sessionkey_gen : start");

    peerkey = EC_KEY_new_by_curve_name(NID_secp521r1);
    if (peerkey == NULL) {
        DUALDAR_LOG("EC_KEY_new_by_curve_name: failed. peerkey == NULL");
        goto exit;
    }
    peergroup = EC_KEY_get0_group(key);
    peerpub_key = EC_POINT_new(peergroup);
    EC_POINT_oct2point(peergroup, peerpub_key, pubkey, pubkey_len, NULL);
    EC_KEY_set_public_key(peerkey, peerpub_key);

    if (!EC_KEY_check_key(peerkey)) {
        DUALDAR_LOG("EC_KEY_check_key for peerkey failed:\n");
    } else {
        DUALDAR_LOG("Public key verified OK for peerkey\n");
    }

    /* Get the peer's public key, and provide the peer with our public key -
     * how this is done will be specific to your circumstances */

    /* Calculate the size of the buffer for the shared secret */
    field_size = EC_GROUP_get_degree(EC_KEY_get0_group(key));
    *secret_len = (field_size+7)/8;

    /* Derive the shared secret */
    *secret_len = ECDH_compute_key(secret, *secret_len, EC_KEY_get0_public_key(peerkey), key, NULL);

    ecdh_keyfree(peerkey);
    DUALDAR_LOG("ecdh_sessionkey_gen : end");

    ret =  DDAR_TZ_API_OK;

exit:
    return ret;
}

/**
 * @Generate random data.
 *
 * @param[in, out] randomBuffer:  The output random data buffer
 * @param[in, out] randomLen: The length of random data to generate.
 * The number of random data bytes generated will be returned in randomLen.
 * */
uint32_t DDAR_TZ_gen_rand_data(
    uint8_t         *randomBuffer,
    uint32_t          *randomLen
)
{
	uint32_t ret= DDAR_TZ_API_OK;
	DUALDAR_LOG("DDAR_TZ_gen_rand_data(SCrypto)");

	if (randomBuffer == NULL) {
		DUALDAR_ERROR("%s invalid input buffer\n", __FUNCTION__);
		return DDAR_TZ_API_ERROR;
	}

	if (randomLen == NULL || *randomLen == 0) {
		DUALDAR_ERROR("%s invalid input buffer length\n", __FUNCTION__);
		return DDAR_TZ_API_ERROR;
	}

	if (*randomLen > MAX_PRNG_LENGTH) {
		DUALDAR_ERROR("%s randomLen too large\n", __FUNCTION__);
		return DDAR_TZ_API_ERROR;
	}

	RAND_bytes(randomBuffer, *randomLen);
    //SPECIALDUMP("randomBuffer", randomBuffer, *randomLen);

	return ret;
}