/*
 * =====================================================================================
 *
 *       Filename:  pebble_core.c
 *
 *    Description:  PEBBLE core functions
 *
 *        Version:  1.0
 *        Created:  06/03/2020
 *       Revision:  none
 *       Compiler:  gcc
 *
 *        Company:  Samsung Electronics
 *        Copyright (c) 2020 by Samsung Electronics, All rights reserved.
 *
 * =====================================================================================
 */

/** Includes */
#include <openssl/rand.h>
#include <openssl/rsa.h>
#include "pebble_core.h"
#include "jwe.h"
#include "jws.h"
#include "cipher.h"
#include "stdbool.h"

//#define LOCAL_TEST
#ifdef LOCAL_TEST
static const unsigned char test_gcm_key[32] = { 0xee, 0xbc, 0x1f, 0x57, 0x48,
		0x7f, 0x51, 0x92, 0x1c, 0x04, 0x65, 0x66, 0x5f, 0x8a, 0xe6, 0xd1, 0x65,
		0x8b, 0xb2, 0x6d, 0xe6, 0xf8, 0xa0, 0x69, 0xa3, 0x52, 0x02, 0x93, 0xa5,
		0x72, 0x07, 0x8f };

static const unsigned char test_gcm_iv[12] = { 0x99, 0xaa, 0x3e, 0x68, 0xed,
		0x81, 0x73, 0xa0, 0xee, 0xd0, 0x66, 0x84 };

static unsigned char test_gcm_pt[32] = { 0xf5, 0x6e, 0x87, 0x05, 0x5b, 0xc3,
		0x2d, 0x0e, 0xeb, 0x31, 0xb2, 0xea, 0xf5, 0x6e, 0x87, 0x05, 0x5b, 0xc3,
		0x2d, 0x0e, 0xeb, 0x31, 0xb2, 0xea, 0xf5, 0x6e, 0x87, 0x05, 0x5b, 0xc3,
		0x2d, 0x0e };
#endif

static unsigned char cert_begin_tag[] = { 0x2D, 0x2D, 0x2D, 0x2D, 0x2D, 0x42,
		0x45, 0x47, 0x49, 0x4E, 0x20, 0x43, 0x45, 0x52, 0x54, 0x49, 0x46, 0x49,
		0x43, 0x41, 0x54, 0x45, 0x2D, 0x2D, 0x2D, 0x2D, 0x2D, 0x0A };
static unsigned char cert_end_tag[] = { 0x0A, 0x2D, 0x2D, 0x2D, 0x2D, 0x2D,
		0x45, 0x4E, 0x44, 0x20, 0x43, 0x45, 0x52, 0x54, 0x49, 0x46, 0x49, 0x43,
		0x41, 0x54, 0x45, 0x2D, 0x2D, 0x2D, 0x2D, 0x2D };

bool drk_ready = false;
uint8_t nonce[ATN_NONCE_MAX_LEN];
uint32_t nonce_len = sizeof(nonce);
bool nonce_ready = false;
char kid[] = "pebble_drk_key"; // a hint for the key
rsa_key_info_t self_key = { 0 }; // self DRK key

/**
 * convert a drk_rsa_private_key_t object to a rsa_key_info_t object
 * input: drk private key, key id (as a hint about the drk key)
 * output: rsa key info
 */
pebble_return_code_t convert_drk_key_to_rsa_info(drk_rsa_private_key_t *drk_key,
		rsa_key_info_t *key, char *kid) {
	if (drk_key->modulus_len > sizeof(key->modulus)
			|| drk_key->priv_expo_len > sizeof(key->priv_expo)
			|| drk_key->pub_expo_len > sizeof(key->pub_expo)
			|| sizeof(kid) > sizeof(key->kid)) {
		PEBBLE_LOG("%s: insufficient buf size in rsa_key_info_t", __FUNCTION__);
		return PEBBLE_KEY_ERROR;
	}
	TEE_MemMove(key->modulus, drk_key->modulus, drk_key->modulus_len);
	key->modulus_len = drk_key->modulus_len;
	TEE_MemMove(key->priv_expo, drk_key->priv_expo, drk_key->priv_expo_len);
	key->priv_expo_len = drk_key->priv_expo_len;
	TEE_MemMove(key->pub_expo, drk_key->pub_expo, drk_key->pub_expo_len);
	key->pub_expo_len = drk_key->pub_expo_len;
	TEE_MemMove(key->kid, kid, sizeof(kid));
	key->kid_len = sizeof(kid);
	return PEBBLE_STATUS_SUCCESS;
}

pebble_return_code_t convert_evp_pkey_to_rsa_info(EVP_PKEY *pkey,
		rsa_key_info_t *key, char *kid) {
	pebble_return_code_t ret = PEBBLE_STATUS_FAIL;
	unsigned char *pn = NULL;
	unsigned char *pe = NULL;
	RSA *rsa = EVP_PKEY_get1_RSA(pkey);
	int n_len = BN_num_bytes(rsa->n);
	int e_len = BN_num_bytes(rsa->e);
	if (n_len > sizeof(key->modulus) || e_len > sizeof(key->pub_expo)
			|| sizeof(kid) > sizeof(key->kid)) {
		PEBBLE_LOG("%s: insufficient buf size in rsa_key_info_t", __FUNCTION__);
		ret = PEBBLE_KEY_ERROR;
		goto exit;
	}
	pn = (unsigned char*) OPENSSL_malloc(n_len);
	pe = (unsigned char*) OPENSSL_malloc(e_len);
	if (!pn || !pe) {
		PEBBLE_LOG("%s: failed to allocate pn and pe", __FUNCTION__);
		ret = PEBBLE_KEY_ERROR;
		goto exit;
	}
	n_len = BN_bn2bin(rsa->n, pn);
	e_len = BN_bn2bin(rsa->e, pe);
	TEE_MemMove(key->modulus, pn, n_len);
	key->modulus_len = n_len;
	TEE_MemMove(key->pub_expo, pe, e_len);
	key->pub_expo_len = e_len;
	TEE_MemMove(key->kid, kid, sizeof(kid));
	key->kid_len = sizeof(kid);
	ret = PEBBLE_STATUS_SUCCESS;
	exit: if (pn != NULL) {
		OPENSSL_cleanse(pn, n_len);
		OPENSSL_free(pn);
	}
	if (pe != NULL) {
		OPENSSL_cleanse(pe, e_len);
		OPENSSL_free(pe);
	}
	return ret;
}

#if 0
/**
 * @brief
 * verify_jws_signature
 * Verify JWS signature
 *
 * @param[in] **pkey                              - enveloped key
 * @param[in]  *tci_msg                           - tci message
 * @param[in]  *size_header_payload_for_signature - length of signature
 *
 * @return PEBBLE status code
 */
static pebble_return_code_t verify_jws_signature(EVP_PKEY **pkey,
		tci_message_t *tci_msg, uint32_t size_header_payload_for_signature) {
	PEBBLE_LOG("verify_jws_signature()");

	pebble_return_code_t ret = PEBBLE_STATUS_SUCCESS;
	EVP_MD_CTX *m_RSAVerifyCtx = EVP_MD_CTX_create();

	if (m_RSAVerifyCtx == NULL) { // SI-16872
		PEBBLE_LOG("EVP_MD_CTX_create FAIL");
		return PEBBLE_STATUS_FAIL;
	}

	if (EVP_DigestVerifyInit(m_RSAVerifyCtx, NULL, EVP_sha256(), NULL, *pkey)
			<= 0) {
		PEBBLE_LOG("EVP_DigestVerifyInit FAIL");
		ret = PEBBLE_STATUS_FAIL;
		goto exit;
	}

	if (EVP_DigestVerifyUpdate(m_RSAVerifyCtx,
			tci_msg->payload.jws_message.data,
			size_header_payload_for_signature)
			<= 0) {
		PEBBLE_LOG("EVP_DigestVerifyUpdate FAIL");
		ret = PEBBLE_STATUS_FAIL;
		goto exit;
	}

	// EVP_DigestVerifyFinal() returns 1 for success; any other value indicates failure.
	ret = EVP_DigestVerifyFinal(m_RSAVerifyCtx, signature.signature,
			signature.signature_len);
	if (!ret) {
		PEBBLE_LOG("EVP_DigestVerifyFinal FAIL");
		PEBBLE_LOG_DEBUG("EVP_DigestVerifyFinal FAIL ret = %d", ret);
		ret = PEBBLE_INVALID_SIGNATURE;
		goto exit;
	}

	PEBBLE_LOG("EVP_DigestVerifyFinal SUCCESS");
	ret = PEBBLE_STATUS_SUCCESS;

	exit: EVP_MD_CTX_cleanup(m_RSAVerifyCtx);
	EVP_MD_CTX_destroy(m_RSAVerifyCtx);

	return ret;
}

/**
 * @brief
 * check_jws_struct
 * Verify and parse PEBBLE JWS policy
 *
 * - Check JWS length
 * - Split base64 JWS
 * - Decode base64 JWS header
 * - Parse JWS header
 * - Validate x509 certificates chains
 * - Decode base64 signature
 * - Parse JWS signature
 * - Validate jws signature
 * - Decode base64 payload
 * - Parse payload
 * - Check DRK certificate
 * - Validate DRKv2 device id
 * - Check policy version
 * - Copy policy to response
 *
 * @param[in] *tci_msg - tci message
 * @return PEBBLE status code
 */
pebble_return_code_t check_jws_struct(tci_message_t *tci_msg)
{
        PEBBLE_LOG("check_jws_struct()");

        pebble_return_code_t ret = PEBBLE_STATUS_SUCCESS;
        tz_pebble_jws_t *jws_data;
        tz_pebble_jws_t *jwsb64_data;
        uint32_t size_header_payload_for_signature;
        EVP_PKEY *pkey = NULL;
        uint8_t *unwrap_object = NULL;
        uint32_t unwrap_object_len = 0;
        uint8_t drk_device_id[DEVICE_ID_B64_LEN] = {0}; // b64( H( H(ID) | H(SERIAL)))

        jws_data = TEE_Malloc(sizeof(tz_pebble_jws_t), 0);
        jwsb64_data = TEE_Malloc(sizeof(tz_pebble_jws_t), 0);

        if (jws_data == NULL || jwsb64_data == NULL)
        {
                ret = PEBBLE_ALLOC_ERROR;
                goto exit;
        }

        PEBBLE_LOG_DEBUG("JWS length = %d", tci_msg->payload.jws_message.len);

        // Check JWS length
        if (tci_msg->payload.jws_message.len <= 0 || tci_msg->payload.jws_message.len > JWS_LEN)
        {
                PEBBLE_LOG("JWS invalid length");
                PEBBLE_LOG_DEBUG("JWS has lenght [ %d ]", tci_msg->payload.jws_message.len);
                ret = PEBBLE_INVALID_JWS;
                goto exit;
        }

        // Split base64 JWS
        ret = split_b64_jws(tci_msg, jwsb64_data, &size_header_payload_for_signature);
        if (ret != PEBBLE_STATUS_SUCCESS)
        {
                PEBBLE_LOG("split_b64_jws FAIL");
                PEBBLE_LOG_DEBUG("split_b64_jws FAIL ret = %d", ret);
                goto exit;
        }

        // Decode base64 JWS header
        PEBBLE_LOG_DEBUG("Decode base64 header");
        jws_data->header_len = HEADER_LEN;
        ret = base64url_decode(jwsb64_data->header, jwsb64_data->header_len,
                               jws_data->header, &jws_data->header_len);
        if (ret != BASE64_OK)
        {
                PEBBLE_LOG("Invalid decode Header");
                PEBBLE_LOG_DEBUG("Invalid decode Header ret = %d", ret);
                ret = PEBBLE_INVALID_JWS;
                goto exit;
        }

        // Parse JWS header
        PEBBLE_LOG_DEBUG("Parse JWS header");
        header.x5c_count = 0;
        ret = fill_jws_msg(jws_data->header, jws_data->header_len, HEADER_TYPE);
        if (ret != PEBBLE_STATUS_SUCCESS)
        {
                PEBBLE_LOG("Parse JWS header FAIL");
                PEBBLE_LOG_DEBUG("Parse JWS header FAIL ret = %d", ret);
                goto exit;
        }

        // Validate x509 certificates chains
        PEBBLE_LOG_DEBUG("Validate cert chain");
        ret = verify_cert_chain(header.x5c, header.x5c_count, &pkey);
        if (ret != PEBBLE_STATUS_SUCCESS)
        {
                PEBBLE_LOG("verify_cert_chain FAIL");
                PEBBLE_LOG_DEBUG("verify_cert_chain FAIL ret = %d", ret);
                goto exit;
        }

        // Decode base64 signature
        PEBBLE_LOG_DEBUG("Decode base64 signature");
        jws_data->signature_len = SIGNATURE_LEN;
        ret = base64url_decode(jwsb64_data->signature, jwsb64_data->signature_len,
                               jws_data->signature, &jws_data->signature_len);
        if (ret != BASE64_OK)
        {
                PEBBLE_LOG("Invalid decode signature");
                PEBBLE_LOG_DEBUG("Invalid decode signature ret = %d", ret);
                ret = PEBBLE_INVALID_JWS;
                goto exit;
        }

        // Parse JWS signature
        PEBBLE_LOG_DEBUG("Parse JWS signature");
        ret = fill_jws_msg_signature(jws_data->signature, jws_data->signature_len);
        if (ret != PEBBLE_STATUS_SUCCESS)
        {
                PEBBLE_LOG("Parse JWS signature FAIL");
                PEBBLE_LOG_DEBUG("Parse JWS signature FAIL ret = %d", ret);
                goto exit;
        }

        // Validate jws signature
        PEBBLE_LOG_DEBUG("Validate JWS signature");
        ret = verify_jws_signature(&pkey, tci_msg, size_header_payload_for_signature);
        if (ret != PEBBLE_STATUS_SUCCESS)
        {
                PEBBLE_LOG("verify_jws_signature FAIL");
                PEBBLE_LOG_DEBUG("verify_jws_signature FAIL ret = %d", ret);
                goto exit;
        }

        // Decode base64 payload
        PEBBLE_LOG_DEBUG("Decode base64 payload");
        jws_data->payload_len = PAYLOAD_LEN;
        ret = base64url_decode(jwsb64_data->payload, jwsb64_data->payload_len,
                               jws_data->payload, &jws_data->payload_len);
        if (ret != BASE64_OK)
        {
                PEBBLE_LOG("Invalid decode payload");
                PEBBLE_LOG_DEBUG("Invalid decode payload ret = %d", ret);
                ret = PEBBLE_INVALID_JWS;
                goto exit;
        }

        // Parse payload
        PEBBLE_LOG_DEBUG("Parse JWS payload - Paylod len %d", jws_data->payload_len);
        ret = fill_jws_msg(jws_data->payload, jws_data->payload_len, PAYLOAD_TYPE);
        if (ret != PEBBLE_STATUS_SUCCESS)
        {
                PEBBLE_LOG("Parse JWS-Payload FAIL");
                PEBBLE_LOG_DEBUG("Parse JWS-Payload FAIL ret = %d", ret);
                goto exit;
        }

        // Check device_id
        ret = check_jws_device_id(payload.device_id, tci_msg->payload.nwd_info);
        if (ret != PEBBLE_STATUS_SUCCESS) {
                PEBBLE_LOG("error on device id check");
                goto exit;
        }

        // Check DRK certificate
        unwrap_object = TEE_Malloc(PEBBLE_DRK_MAX_BUF_LEN, 0);

        //SRR-14493, SRR-14486
        if(unwrap_object == NULL)
        {
                ret = PEBBLE_ALLOC_ERROR;
                PEBBLE_LOG("Error to alloc memory");
                goto exit;
        }

        unwrap_object_len = PEBBLE_DRK_MAX_BUF_LEN;
        uint32_t wrap_len = tci_msg->payload.nwd_info.pebble_key_len;

        ret = unwrap(tci_msg->payload.nwd_info.pebble_key, &wrap_len, tci_msg->payload.nwd_info.is_wrapped_key, unwrap_object, &unwrap_object_len);
        tci_msg->payload.nwd_info.pebble_key_len = wrap_len;

        if (ret != PEBBLE_STATUS_SUCCESS)
        {
                PEBBLE_LOG("DRK unwrap failed");
                goto exit;
        }

        ret = get_b64_hash(unwrap_object, unwrap_object_len, (char *)drk_device_id);
        if (ret != PEBBLE_STATUS_SUCCESS)
        {
                PEBBLE_LOG("Error DRK unwrap");
                PEBBLE_LOG_DEBUG("DRK unwrap ret = %d", ret);
                goto exit;
        }
        PEBBLE_LOG_DEBUG("drk_device_id = %s", drk_device_id);

        // Validate DRKv2 device id
        ret = validate_drk_device_id(drk_device_id, tci_msg->payload.nwd_info);
        if (ret != PEBBLE_STATUS_SUCCESS)
        {
                PEBBLE_LOG("DRK device id mismatch");
                goto exit;
        }
        PEBBLE_LOG("Device id match!");

        // Copy policy to response
        PEBBLE_LOG_DEBUG("Policy Value = %d", payload.device_block);
        //tci_msg->payload.jws_message.policy_value = payload.device_block;

exit:
        if (unwrap_object != NULL)
        {
                TEE_Free(unwrap_object);
        }
        if (pkey != NULL) {
                EVP_PKEY_free(pkey);
        }
        TEE_Free(jwsb64_data);
        TEE_Free(jws_data);

        return ret;
}
#endif

/**
 * @brief
 * pebble_ICCC_check
 *
 * Verify device integrity with ICCC device status API
 * @return PEBBLE status code
 */
pebble_return_code_t pebble_ICCC_check(void) {
	PEBBLE_LOG("pebble_ICCC_check()");
	uint32_t result_code = ICCC_STATUS_RETURN_NOT_SECURE;
	uint32_t ret = ICCC_ERROR_DEVICE_STATUS_FAILED;
	uint32_t comp_type = ICCC_STATUS_COMP_TYPE_SOFT_INTEGRITY;
	uint32_t ta_status_msg_len = 0;
	uint8_t ta_status_msg[ICCC_STATUS_MAX_RESULT_MESSAGE];

	ret = Iccc_DeviceStatus_TA(comp_type, ta_status_msg, sizeof(ta_status_msg),
			&ta_status_msg_len, &result_code);
	if (ret != ICCC_SUCCESS) {
		PEBBLE_LOG("Unable to verify device status");
		PEBBLE_LOG_DEBUG(
				"Unable to verify device status - Iccc_DeviceStatus_TA failed");
		return PEBBLE_DEVICE_COMPROMISED;
	}

	if (result_code != ICCC_STATUS_RETURN_SECURE) {
		PEBBLE_LOG("Device is not secure");
		PEBBLE_LOG_DEBUG(
				"Iccc_DeviceStatus_TA return is not ICCC_STATUS_RETURN_SECURE. Device is not secure ");
		return PEBBLE_DEVICE_COMPROMISED;
	}

	return PEBBLE_DEVICE_OK;
}

/**
 * @brief
 * pebble_ICCC_check_em_status
 *
 * Verify device EM Status with ICCC read data API
 * - em_status = 0x0 -> PROD device
 * - em_status = 0x1 -> DEV  device
 * - em_status = 0x2 -> Eng token is installed on PROD device
 *
 * @param[out] em_status - EM Status value
 * @return PEBBLE Status
 */
pebble_return_code_t pebble_ICCC_check_em_status(uint32_t *em_status) {
	PEBBLE_LOG("pebble_ICCC_check_em_status()");
	uint32_t ret = ICCC_ERROR_READ_FAILED;

	ret = Iccc_ReadData_TA(EM_STATUS, em_status);
	if (ret != ICCC_SUCCESS) {
		PEBBLE_LOG("Unable to read em_status");
		PEBBLE_LOG_DEBUG(
				"Unable to read em_status - Iccc_ReadData_TA failed %d", ret);
		return PEBBLE_READ_EM_STATUS_FAIL;
	}

	return PEBBLE_STATUS_SUCCESS;
}

/**
 * @brief
 * load key and cert chain for pebble TA
 *
 * @param[in] *tci_msg - tci message
 * @return PEBBLE status code
 */
pebble_return_code_t load_cert(tci_message_t *cmd, tci_message_t *resp) {
	if (cmd->payload.load_cert_cmd.wrapped_key.len > PEBBLE_DRK_MAX_BUF_LEN) {
		PEBBLE_LOG("%s: len from user input exceeds the limit", __FUNCTION__);
		return PEBBLE_INVALID_USER_INPUT_LEN;
	}
	uint32_t ret = PEBBLE_GEN_RESPONSE_FAIL;
	uint8_t *unwrap_object = NULL;
	uint32_t unwrap_object_len = 0;
	int i = 0;
	drk_parsed_object_t drk_data;

	PEBBLE_LOG("load_cert()");

	// Check DRK certificate
	unwrap_object = TEE_Malloc(PEBBLE_DRK_MAX_BUF_LEN, 0);

	if (unwrap_object == NULL) {
		ret = PEBBLE_ALLOC_ERROR;
		PEBBLE_LOG("Error to alloc memory");
		goto exit;
	}

	unwrap_object_len = PEBBLE_DRK_MAX_BUF_LEN;
	uint32_t wrap_len = cmd->payload.load_cert_cmd.wrapped_key.len;

	ret = unwrap(cmd->payload.load_cert_cmd.wrapped_key.buf, &wrap_len, 0,
			unwrap_object, &unwrap_object_len);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG("DRK unwrap failed");
		goto exit;
	}

	// Get certificates chain RSA key
	ret = get_cert_chain_rsakey(unwrap_object, unwrap_object_len, &drk_data);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG("get_cert_chain_rsakey FAIL");
		PEBBLE_LOG_DEBUG("get_cert_chain_rsakey FAIL ret = %d", ret);
		goto exit;
	}

	// Convert DER to PEM
	ret = convert_der_to_b64(drk_data.drk_cert_chain.cert,
			drk_data.drk_cert_chain.num_certs);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG("convert_der_to_b64 FAIL");
		PEBBLE_LOG_DEBUG("convert_der_to_b64 FAIL ret = %d", ret);
		goto exit;
	}

	for (i = 0; i < drk_data.drk_cert_chain.num_certs; i++) {
		int len = 0;
		TEE_MemMove(resp->payload.load_cert_resp.cert_chain.cert[i].blob,
				cert_begin_tag, sizeof(cert_begin_tag));
		len += sizeof(cert_begin_tag);
		TEE_MemMove(resp->payload.load_cert_resp.cert_chain.cert[i].blob + len,
				drk_data.drk_cert_chain.cert[i].blob,
				drk_data.drk_cert_chain.cert[i].len);
		len += drk_data.drk_cert_chain.cert[i].len;
		TEE_MemMove(resp->payload.load_cert_resp.cert_chain.cert[i].blob + len,
				cert_end_tag, sizeof(cert_end_tag));
		len += sizeof(cert_end_tag);
		resp->payload.load_cert_resp.cert_chain.cert[i].len = len;
	}
	resp->payload.load_cert_resp.cert_chain.num_certs =
			drk_data.drk_cert_chain.num_certs;
	convert_drk_key_to_rsa_info(&drk_data.drk_rsa_private_key, &self_key, kid);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG("%s: convert_drk_key_to_rsa_info failed", __FUNCTION__);
		goto exit;
	}
	drk_ready = true;

	exit: if (unwrap_object != NULL) {
		memset(unwrap_object, 0, PEBBLE_DRK_MAX_BUF_LEN);
		TEE_Free(unwrap_object);
	}
	memset(&drk_data, 0, sizeof(drk_data));

	return ret;
}

/*
 * encrypt given plaintext by RSA OAEPPadding md=SHA256, mgf=mfg1, mgf-md=SHA1.
 * returns the size of the encrypted data (i.e., RSA_size(rsa)).
 * Purpose: mainly used to wrap an AES key with Android KeyStore public key.
 * This function is modified based on openssl's RSA_eay_public_encrypt function.
 */
int RSA_public_encrypt_Pebble(const unsigned char *from, int flen,
		EVP_PKEY *pkey, unsigned char *to) //const char *n_str
{
	int padding = RSA_PKCS1_OAEP_PADDING;
	BIGNUM *f, *ct;
	int i, j, k, num = 0, r = -1;
	unsigned char *buf = NULL;
	BN_CTX *ctx = NULL;
	int ret;
	RSA *rsa = EVP_PKEY_get0_RSA(pkey);

//#ifdef OPENSSL_FIPS
//    ret = FIPS_mode_set(1 /*on*/);
//    if(ret != 1)
//    {
//    	PEBBLE_LOG("failed to setup FIPS mode");
//    }
//
//    if (FIPS_mode() && !(rsa->flags & RSA_FLAG_NON_FIPS_ALLOW)
//        && (BN_num_bits(rsa->n) < OPENSSL_RSA_FIPS_MIN_MODULUS_BITS))
//    {
//        PEBBLE_LOG("RSA_R_KEY_SIZE_TOO_SMALL");
//        return -1;
//    }
//#endif

//	if (BN_num_bits(rsa->n) > OPENSSL_RSA_MAX_MODULUS_BITS) {
//		PEBBLE_LOG("RSA_R_MODULUS_TOO_LARGE");
//		return -1;
//	}

	if (BN_ucmp(rsa->n, rsa->e) <= 0) {
		PEBBLE_LOG("RSA_R_BAD_E_VALUE");
		return -1;
	}

//	/* for large moduli, enforce exponent limit */
//	if (BN_num_bits(rsa->n) > OPENSSL_RSA_SMALL_MODULUS_BITS) {
//		if (BN_num_bits(rsa->n) > OPENSSL_RSA_MAX_PUBEXP_BITS) {
//			PEBBLE_LOG("RSA_R_BAD_E_VALUE");
//			return -1;
//		}
//	}

	if ((ctx = BN_CTX_new()) == NULL)
		goto err;
	BN_CTX_start(ctx);
	f = BN_CTX_get(ctx);
	ct = BN_CTX_get(ctx);
	num = BN_num_bytes(rsa->n); //2048bits/8 = 256 bytes
	buf = (unsigned char*) OPENSSL_malloc(num);
	if (!f || !ct || !buf) {
		PEBBLE_LOG("ERR_R_MALLOC_FAILURE");
		goto err;
	}
	// pad the plaintext
	//i = RSA_padding_add_PKCS1_OAEP_SHA256(buf, num, from, flen, NULL, 0);
	i = RSA_padding_add_PKCS1_OAEP_mgf1(buf, num, from, flen,
	NULL, 0, EVP_sha256(), EVP_sha1()); // this line is changed by Jianwei

	if (i <= 0)
		goto err;

	if (BN_bin2bn(buf, num, f) == NULL)
		goto err;

	if (BN_ucmp(f, rsa->n) >= 0) {
		/* usually the padding functions would catch this */
		PEBBLE_LOG("RSA_R_DATA_TOO_LARGE_FOR_MODULUS");
		goto err;
	}
//    if (rsa->flags & RSA_FLAG_CACHE_PUBLIC) // commented because openssl does not allow accessing private fields like rsa->_method_mod_n
//        if (!BN_MONT_CTX_set_locked(&rsa->_method_mod_n, CRYPTO_LOCK_RSA, rsa->n, ctx))
//            goto err;
	// encrypt the data f
	if (!BN_mod_exp(ct, f, rsa->e, rsa->n, ctx)) //RSA_get0_e(rsa)
		goto err;

	/* put in leading 0 bytes if the number is less than the length of the modulus */
	j = BN_num_bytes(ct);  // ret is the encrypted data
	i = BN_bn2bin(ct, &(to[num - j]));
	for (k = 0; k < (num - i); k++)
		to[k] = 0;
	r = num;
	RSA_free(rsa);
	err: if (ctx != NULL) {
		BN_CTX_end(ctx);
		BN_CTX_free(ctx);
	}
	if (buf != NULL) {
		OPENSSL_cleanse(buf, num);
		OPENSSL_free(buf);
	}
	return r;
}

///*
// * AES encryption with GCM mode.
// * Input: plaintext pt, key, initialization vector iv, additional authenticated data aad
// * Output: ciphertext ct, authentication tag
// * Purpose: mainly used to wrap a secret key with a temporary AES key.
// */
//int aes_gcm_encrypt(unsigned char *gcm_pt, int gcm_pt_len,
//		unsigned char gcm_key[WRAPPING_AES_KEY_LEN],
//		unsigned char gcm_iv[WRAPPING_AES_IV_LEN],
//		unsigned char *gcm_aad,
//		int gcm_aad_len,
//		unsigned char *gcm_ct,
//		unsigned char gcm_tag[WRAPPING_AES_TAG_LEN]) {
//
//    EVP_CIPHER_CTX ctx;
//    int tmp_len;
//    int tmp_ciphertext_len;
//    int gcm_ct_len;
//    uint32_t ret = PEBBLE_STATUS_FAIL;
//
//    /* 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)) {
//        PEBBLE_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, WRAPPING_AES_IV_LEN, NULL)){
//        PEBBLE_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, gcm_key, gcm_iv)) {
//        PEBBLE_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, gcm_aad, gcm_aad_len)) {
//        PEBBLE_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, gcm_ct, &tmp_len, gcm_pt, gcm_pt_len)) {
//        PEBBLE_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, gcm_ct + tmp_len, &tmp_len)) {
//        PEBBLE_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, WRAPPING_AES_TAG_LEN, gcm_tag)) {
//        PEBBLE_LOG("%s: EVP_CIPHER_CTX_ctrl(EVP_CTRL_GCM_GET_TAG, tag) failed", __FUNCTION__);
//        goto exit;
//    }
//
//    gcm_ct_len = tmp_ciphertext_len;
//
//    ret = PEBBLE_STATUS_SUCCESS;
//
//exit:
//    EVP_CIPHER_CTX_cleanup(&ctx);
//    return ret;
//}

/*
 * encrypt the given secret with a random AES key and wrap the AES key with the provided public key
 * Input: secret, AES-IV (gcm_iv), AES-AAD (gcm_aad), public key (pkey),
 * Output: encrypted secret (secret_ct), initial vector (gcm_iv), authenticated tag (gcm_tag), wrapped AES key (wrapped_gcm_key)
 * Purpose: wrap the secret key with a public key from Android KeyStore so that
 * it can be securely imported to KeyStore in normal world.
 *
 */
pebble_return_code_t two_step_wrap_with_android_keystore_public_key(
		unsigned char *secret, int secret_len, EVP_PKEY *pkey, uint8_t *gcm_aad,
		uint32_t gcm_aad_len, uint8_t *secret_ct, uint32_t *secret_ct_len,
		uint8_t *gcm_iv, uint32_t gcm_iv_len, uint8_t *gcm_tag,
		uint32_t gcm_tag_len, uint8_t *wrapped_gcm_key,
		uint32_t *wrapped_gcm_key_len) {
	int ret, ctlen;

	// step 1: use a random aes-gcm key to encrypt the secret
	unsigned char gcm_key[WRAPPING_AES_KEY_LEN]; //used to encrypt master key, cannot be leaked
	unsigned char gcm_iv_buf[gcm_iv_len];
#ifdef LOCAL_TEST
	TEE_MemMove(gcm_key, (void*) test_gcm_key, WRAPPING_AES_KEY_LEN);
	TEE_MemMove(gcm_iv_buf, (void*) test_gcm_iv, WRAPPING_AES_IV_LEN);
#else        
	RAND_bytes((unsigned char*) gcm_key, WRAPPING_AES_KEY_LEN);
	RAND_bytes((unsigned char*) gcm_iv_buf, gcm_iv_len);
#endif
	*secret_ct_len = secret_len;
//	ret = aes_gcm_encrypt((unsigned char*) secret, secret_len, gcm_key,
//			gcm_iv_buf, gcm_aad, gcm_aad_len, secret_ct, gcm_tag);
	ret = pebble_aes_gcm_encrypt(gcm_key, WRAPPING_AES_KEY_LEN, gcm_iv_buf,
			gcm_iv_len, gcm_aad, gcm_aad_len, secret, secret_len, secret_ct,
			gcm_tag);
	TEE_MemMove(gcm_iv, gcm_iv_buf, sizeof(gcm_iv_buf));
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG(
				"two_step_wrap_with_android_keystore_public_key - failed to encrypt the secret with AES key");
		return PEBBLE_WRAP_FAILED;
	}

	// step 2: use the provided KeyStore public key to encrypt the aes key
	ctlen = RSA_public_encrypt_Pebble((const unsigned char*) gcm_key,
			sizeof(gcm_key), pkey, wrapped_gcm_key);
	if (ctlen == 0) {
		PEBBLE_LOG(
				"two_step_wrap_with_android_keystore_public_key - RSA_public_encrypt_Pebble failed");
		return PEBBLE_WRAP_FAILED;
	}

	return PEBBLE_STATUS_SUCCESS;
}

/*
 * generate a random 256 bit master key, (1) encrypt it with TA, (2) wrap it with Android KeyStore public key
 * Input: cmd (type gen_master_key_cmd_t) contains AES-GCM AAD(gcm_aad), Android Key Store certificate chain.
 * Output: resp (type gen_master_key_resp_t) contains master key wrapped by TA, master key encrypted by AES-GCM (secret_ct, gcm_iv, gcm_tag), aes key encrypted (wrapped_gcm_key) by provided public key.
 */
pebble_return_code_t generate_master_key_and_wrap_it(tci_message_t *cmd,
		tci_message_t *resp) {
	if (cmd->payload.gen_master_key_cmd.gcm_aad.len > WRAPPING_AES_AAD_LEN) {
		PEBBLE_LOG("%s: len from user input exceeds the limit", __FUNCTION__);
		return PEBBLE_INVALID_USER_INPUT_LEN;
	}
	unsigned char master_key[SK_MASTER_KEY_LEN];
	pebble_return_code_t ret = PEBBLE_STATUS_FAIL;
	EVP_PKEY *pkey = NULL;
	uint8_t *key_wrapped_ta;
	uint32_t key_wrapped_ta_len;
	uint8_t *secret_ct;
	uint32_t secret_ct_len;
	uint8_t *wrapped_gcm_key;
	uint32_t wrapped_gcm_key_len;

	// Validate x509 certificates chains
	PEBBLE_LOG_DEBUG("Validate cert chain");
#ifdef LOCAL_TEST
	const char *test_nonce = "hello world";
	ret = verify_aks_cert_chain(cmd->payload.gen_master_key_cmd.cert_chain,
			&pkey, (uint8_t*) test_nonce, strnlen(test_nonce, ATN_NONCE_MAX_LEN));
#else
	if (nonce_ready)
        ret = verify_aks_cert_chain(cmd->payload.gen_master_key_cmd.cert_chain, &pkey, nonce, nonce_len);
	else {
		PEBBLE_LOG("%s: attestation nonce is not ready. Please call CMD_GEN_ATN_NONCE first", __FUNCTION__);
		return PEBBLE_NONCE_NOT_READY; //return this error to NWD, then NWD should call CMD_GEN_ATN_NONCE
	}
#endif
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG("verify_aks_cert_chain FAIL");
		PEBBLE_LOG_DEBUG("verify_aks_cert_chain FAIL ret = %d", ret);
		return PEBBLE_WRAP_FAILED;
	}
	//PEBBLE_LOG("verify_aks_cert_chain success");

	// generate a random master key
#ifdef LOCAL_TEST
	TEE_MemMove(master_key, test_gcm_pt, SK_MASTER_KEY_LEN);
#else
	RAND_bytes((unsigned char*) master_key, SK_MASTER_KEY_LEN);
#endif   
//	printLongRangeHex(master_key, SK_MASTER_KEY_LEN, "master_key before wrap");

	key_wrapped_ta = TEE_Malloc(SK_WRAP_KEY_LEN, 0);
	secret_ct = TEE_Malloc(SK_WRAP_KEY_LEN, 0);
	wrapped_gcm_key = TEE_Malloc(WRAPPING_RSA_N_BYTES_LEN, 0);

	if (key_wrapped_ta == NULL || secret_ct == NULL || wrapped_gcm_key == NULL) {
		ret = PEBBLE_ALLOC_ERROR;
		PEBBLE_LOG("Error to alloc memory");
		goto exit;
	}
	key_wrapped_ta_len = SK_WRAP_KEY_LEN;
	secret_ct_len = SK_WRAP_KEY_LEN;
	wrapped_gcm_key_len = WRAPPING_RSA_N_BYTES_LEN;

	// 1. wrap it with TA
	ret = wrap((uint8_t*) master_key, sizeof(master_key), key_wrapped_ta,
			&key_wrapped_ta_len);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG(
				"generate_master_key_and_wrap_it - failed to wrap the master key with TA");
		ret = PEBBLE_WRAP_FAILED;  // need to add to tz_pebble_interface.h
		goto exit;
	}
	//PEBBLE_LOG("wrap done");

	// 2. wrap it with Android KeyStore public key
	ret = two_step_wrap_with_android_keystore_public_key(master_key,
			sizeof(master_key), pkey,
			cmd->payload.gen_master_key_cmd.gcm_aad.buf,
			cmd->payload.gen_master_key_cmd.gcm_aad.len, secret_ct,
			&secret_ct_len, resp->payload.gen_master_key_resp.gcm_iv,
			WRAPPING_AES_IV_LEN, resp->payload.gen_master_key_resp.gcm_tag,
			WRAPPING_AES_TAG_LEN, wrapped_gcm_key, &wrapped_gcm_key_len);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG(
				"generate_master_key_and_wrap_it - failed to wrap the master key with Android KeyStore public key");
		ret = PEBBLE_WRAP_FAILED;  // need to add to tz_pebble_interface.h
		goto exit;
	}
//	printLongRangeHex(master_key, sizeof(master_key), "master_key");
//	printLongRangeHex(key_wrapped_ta, key_wrapped_ta_len, "key_wrapped_ta");
//	printLongRangeHex(secret_ct, secret_ct_len, "secret_ct");
//	printLongRangeHex(wrapped_gcm_key, wrapped_gcm_key_len, "wrapped_gcm_key");

	resp->payload.gen_master_key_resp.master_key_wrapped_ta.len =
			key_wrapped_ta_len;
	TEE_MemMove(resp->payload.gen_master_key_resp.master_key_wrapped_ta.buf,
			key_wrapped_ta, key_wrapped_ta_len);
	resp->payload.gen_master_key_resp.secret_ct.len = secret_ct_len;
	TEE_MemMove(resp->payload.gen_master_key_resp.secret_ct.buf, secret_ct,
			secret_ct_len);
	resp->payload.gen_master_key_resp.wrapped_gcm_key.len = wrapped_gcm_key_len;
	TEE_MemMove(resp->payload.gen_master_key_resp.wrapped_gcm_key.buf,
			wrapped_gcm_key, wrapped_gcm_key_len);

	ret = PEBBLE_STATUS_SUCCESS;

	exit: if (key_wrapped_ta != NULL) {
		TEE_Free(key_wrapped_ta);
	}
	if (secret_ct != NULL) {
		TEE_Free(secret_ct);
	}
	if (wrapped_gcm_key != NULL) {
		TEE_Free(wrapped_gcm_key);
	}
	memset(master_key, 0, sizeof(master_key));
	return ret;
}

/*
 * derive a 256 bit sub key from given master key
 * Input: cmd (type derive_subkey_cmd_t) contains wrapped master key, subkey alias, gcm aad, certificate chain
 * Output: resp (type derive_subkey_resp_t) contains subkey encrypted by AES-GCM (secret_ct, gcm_iv, gcm_tag), aes key encrypted (wrapped_gcm_key) by provided public key.
 */
pebble_return_code_t derive_subkey_and_wrap_it(tci_message_t *cmd,
		tci_message_t *resp) {
	if (cmd->payload.derive_subkey_cmd.master_key_wrapped_ta.len > SK_WRAP_KEY_LEN ||
			cmd->payload.derive_subkey_cmd.alias.len > KEY_ALIAS_LEN ||
			cmd->payload.derive_subkey_cmd.gcm_aad.len > WRAPPING_AES_AAD_LEN) {
		PEBBLE_LOG("%s: len from user input exceeds the limit", __FUNCTION__);
		return PEBBLE_INVALID_USER_INPUT_LEN;
	}
	int r;
	pebble_return_code_t ret;
	uint8_t master_key[SK_WRAP_KEY_LEN];
	uint32_t master_key_len = sizeof(master_key);
	unsigned char subkey[SK_SUBKEY_LEN];
	EVP_PKEY *pkey = NULL;
	uint8_t *key_wrapped_ta;
	uint32_t key_wrapped_ta_len;
	uint8_t *secret_ct;
	uint32_t secret_ct_len;
	uint8_t *wrapped_gcm_key;
	uint32_t wrapped_gcm_key_len;

	// Validate x509 certificates chains
	PEBBLE_LOG_DEBUG("Validate cert chain");
#ifdef LOCAL_TEST
	const char *test_nonce = "hello world";
	ret = verify_aks_cert_chain(cmd->payload.derive_subkey_cmd.cert_chain,
			&pkey, (uint8_t*) test_nonce, strnlen(test_nonce, ATN_NONCE_MAX_LEN));
#else
	if (nonce_ready)
        ret = verify_aks_cert_chain(cmd->payload.gen_master_key_cmd.cert_chain, &pkey, nonce, nonce_len);
	else {
		PEBBLE_LOG("%s: attestation nonce is not ready. Please call CMD_GEN_ATN_NONCE first", __FUNCTION__);
		return PEBBLE_NONCE_NOT_READY; //return this error to NWD, then NWD should call CMD_GEN_ATN_NONCE
	}
#endif
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG("verify_aks_cert_chain FAIL");
		PEBBLE_LOG_DEBUG("verify_aks_cert_chain FAIL ret = %d", ret);
		return PEBBLE_WRAP_FAILED;
	}
	//PEBBLE_LOG("verify_aks_cert_chain success");

	key_wrapped_ta = TEE_Malloc(SK_WRAP_KEY_LEN, 0);
	secret_ct = TEE_Malloc(SK_WRAP_KEY_LEN, 0);
	wrapped_gcm_key = TEE_Malloc(WRAPPING_RSA_N_BYTES_LEN, 0);

	if (key_wrapped_ta == NULL || secret_ct == NULL || wrapped_gcm_key == NULL) {
		ret = PEBBLE_ALLOC_ERROR;
		PEBBLE_LOG("Error to alloc memory");
		goto exit;
	}
	key_wrapped_ta_len = SK_WRAP_KEY_LEN;
	secret_ct_len = SK_WRAP_KEY_LEN;
	wrapped_gcm_key_len = WRAPPING_RSA_N_BYTES_LEN;

	TEE_MemMove(key_wrapped_ta,
			cmd->payload.derive_subkey_cmd.master_key_wrapped_ta.buf,
			cmd->payload.derive_subkey_cmd.master_key_wrapped_ta.len);
	key_wrapped_ta_len =
			cmd->payload.derive_subkey_cmd.master_key_wrapped_ta.len;
//	printLongRangeHex(key_wrapped_ta, key_wrapped_ta_len, "key_wrapped_ta");

	// unwrap the provided wrapped master key
	ret = unwrap(key_wrapped_ta, &key_wrapped_ta_len, 1, (uint8_t*) master_key,
			&master_key_len);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG(
				"derive_subkey_and_wrap_it - failed to unwrap the provided master key");
		return PEBBLE_UNWRAP_FAILED;
	}
	if (master_key_len != SK_MASTER_KEY_LEN) {
		PEBBLE_LOG(
				"derive_subkey_and_wrap_it - provided master key's length is not %d bytes",
				SK_MASTER_KEY_LEN);
		return PEBBLE_UNWRAP_FAILED;
	}

	// derive subkey from master key and subkey alias
	r = PKCS5_PBKDF2_HMAC_SHA1(
			(char*) (cmd->payload.derive_subkey_cmd.alias.buf),
			cmd->payload.derive_subkey_cmd.alias.len,
			(const unsigned char*) master_key,
			SK_MASTER_KEY_LEN, 1000, SK_SUBKEY_LEN, subkey); // use PKCS5_PBKDF2_HMAC(...) to choose hashing function
	if (r == 0) {
		PEBBLE_LOG(
				"derive_subkey_and_wrap_it - PKCS5_PBKDF2_HMAC_SHA1 deriving subkey failed");
		return PEBBLE_WRAP_FAILED;
	}

	// wrap subkey with Android KeyStore public key
	ret = two_step_wrap_with_android_keystore_public_key(subkey, sizeof(subkey),
			pkey, cmd->payload.derive_subkey_cmd.gcm_aad.buf,
			cmd->payload.derive_subkey_cmd.gcm_aad.len, secret_ct,
			&secret_ct_len, resp->payload.derive_subkey_resp.gcm_iv,
			WRAPPING_AES_IV_LEN, resp->payload.derive_subkey_resp.gcm_tag,
			WRAPPING_AES_TAG_LEN, wrapped_gcm_key, &wrapped_gcm_key_len);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG(
				"derive_subkey_and_wrap_it - failed to wrap the subkey with Android KeyStore public key");
		ret = PEBBLE_WRAP_FAILED;  // need to add to tz_pebble_interface.h
		goto exit;
	}
//	printLongRangeHex(subkey, sizeof(subkey), "subkey");
//	printLongRangeHex(secret_ct, secret_ct_len, "secret_ct");
//	printLongRangeHex(wrapped_gcm_key, wrapped_gcm_key_len, "wrapped_gcm_key");

	resp->payload.derive_subkey_resp.secret_ct.len = secret_ct_len;
	TEE_MemMove(resp->payload.derive_subkey_resp.secret_ct.buf, secret_ct,
			secret_ct_len);
	resp->payload.derive_subkey_resp.wrapped_gcm_key.len = wrapped_gcm_key_len;
	TEE_MemMove(resp->payload.derive_subkey_resp.wrapped_gcm_key.buf,
			wrapped_gcm_key, wrapped_gcm_key_len);

	ret = PEBBLE_STATUS_SUCCESS;
	exit: if (key_wrapped_ta != NULL) {
		TEE_Free(key_wrapped_ta);
	}

	if (secret_ct != NULL) {
		TEE_Free(secret_ct);
	}
	if (wrapped_gcm_key != NULL) {
		TEE_Free(wrapped_gcm_key);
	}
	memset(master_key, 0, sizeof(master_key));
	memset(subkey, 0, sizeof(subkey));

	return ret;
}

/*
 * generate challenge (base64 encoded 16 random bytes) for Android Key Store certificate
 * Input: cmd (type gen_atn_nonce_cmd_t) contains len (required challenge length which is not used now)
 * Output: resp (type gen_atn_nonce_resp_t) contains nonce (base64 encoded 16 random bytes)
 */
pebble_return_code_t gen_atn_nonce(tci_message_t *cmd, tci_message_t *resp) {
//	if (cmd->payload.gen_atn_nonce_cmd.len > ATN_NONCE_MAX_LEN || cmd->payload.gen_atn_nonce_cmd.len < ATN_NONCE_MIN_LEN) {
//		PEBBLE_LOG("%s: len from user input exceeds the limit", __FUNCTION__);
//		return PEBBLE_INVALID_USER_INPUT_LEN;
//	}
	uint8_t nonce_bin[ATN_NONCE_MAX_LEN / 2];
	if (nonce_ready) {
		PEBBLE_LOG("nonce is already generated, will replace it");
	}
	if (RAND_bytes(nonce_bin, sizeof(nonce_bin)) == 0) {
		PEBBLE_LOG("failed to generate nonce");
		return PEBBLE_GEN_NONCE_FAIL;
	}
	nonce_len = ATN_NONCE_MAX_LEN;
	if (base64_encode(nonce_bin, sizeof(nonce_bin), nonce,
			&nonce_len) != BASE64_OK) {
		PEBBLE_LOG("failed to encode nonce");
		nonce_ready = false;
		return PEBBLE_GEN_NONCE_FAIL;
	}
	resp->payload.gen_atn_nonce_resp.nonce.len = nonce_len;
	TEE_MemMove(resp->payload.gen_atn_nonce_resp.nonce.buf, nonce, nonce_len);
	nonce_ready = true;

	memset(nonce_bin, 0, sizeof(nonce_bin));
	return PEBBLE_STATUS_SUCCESS;
}

/**
 *  Input: wrapped master key, buddies' cert chains
 *  Output: a backup request jws string for each buddy
 */
pebble_return_code_t gen_backup_req(tci_message_t *cmd, tci_message_t *resp) {
	if (cmd->payload.gen_backup_req_cmd.wrapped_master_key.len > SK_WRAP_KEY_LEN
			|| cmd->payload.gen_backup_req_cmd.k > MAX_BUDDIES) {
		PEBBLE_LOG("%s: len from user input exceeds the limit", __FUNCTION__);
		return PEBBLE_INVALID_USER_INPUT_LEN;
	}
	pebble_return_code_t ret;
	if (!drk_ready) {
		ret = PEBBLE_DRK_NOT_READY;
		goto exit;
	}
	uint32_t master_key_len, payload_len, r;
	uint8_t master_key[SK_WRAP_KEY_LEN];
	uint8_t shares[MAX_BUDDIES][SSS_SHARE_MAX_LEN],
			checksums[MAX_BUDDIES][CHECKSUM_LEN];
	size_t share_len = 0, meta_len = 0;
	uint32_t wrapped_master_key_len =
			cmd->payload.gen_backup_req_cmd.wrapped_master_key.len;
	uint8_t n = cmd->payload.gen_backup_req_cmd.n;
	uint8_t k = cmd->payload.gen_backup_req_cmd.k;
	rsa_key_info_t buddy_key = { 0 };
	uint32_t jwe_len, jws_len;
	uint8_t jwe_out[JWE_JWS_MAX_SIZE], jws_out[JWE_JWS_MAX_SIZE-1];

	if (n > MAX_BUDDIES || n < k || k < 1) {
		PEBBLE_LOG("%s: invalid user input. Ensure 1 <= k <= n <= %d", __FUNCTION__, MAX_BUDDIES);
		ret = PEBBLE_INVALID_USER_INPUT_LEN;
		goto exit;
	}

	// unwrap the provided wrapped master key
	master_key_len = sizeof(master_key);
	ret = unwrap(cmd->payload.gen_backup_req_cmd.wrapped_master_key.buf,
			&wrapped_master_key_len, 1, (uint8_t*) master_key, &master_key_len);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG("%s: failed to unwrap the provided master key", __FUNCTION__);
		goto exit;
	}
	if (master_key_len != SK_MASTER_KEY_LEN) {
		PEBBLE_LOG("%s: provided master key's length is %d, not %d",
				__FUNCTION__, master_key_len, SK_MASTER_KEY_LEN);
		ret = PEBBLE_UNWRAP_FAILED;
		goto exit;
	}
	// split master key into n shares
	split((uint8_t*) master_key, SK_MASTER_KEY_LEN, n, k, shares, &share_len);
	if (share_len != SK_MASTER_KEY_LEN + 1) {
		PEBBLE_LOG("%s: expect share len is %d, but get %d", __FUNCTION__, SK_MASTER_KEY_LEN + 1, share_len);
		ret = PEBBLE_SSS_SPLIT_ERROR;
		goto exit;
	}
	// for each share, generate a checksum
	for (int i = 0; i < n; i++) {
		r = gen_checksum(shares[i], share_len, checksums[i]); // generate checksum
		if (r != 1) {
			PEBBLE_LOG("%s: failed to generate checksum for key share",
					__FUNCTION__);
			ret = PEBBLE_GEN_CHECKSUM_FAIL;
			goto exit;
		}
	}
	/** for each share, create a payload (share + n checksums) using customized format: (k is contained in metadata)
	 *  "[share][share number n][share index i][all shares’ checksums concatenated][metadata_json_for_NWD_use]"
	 *     33 B      1 B             1 B            32*n B                             len doesn't matter
	 *  Then, encrypt it with JWE and sign it with JWS
	 */
	uint8_t payloads[MAX_BUDDIES][SHARE_PAYLOAD_LEN];
	for (uint8_t i = 0; i < n; i++) {
		TEE_MemMove(payloads[i], shares[i], SSS_SHARE_MAX_LEN); // share
		payload_len = SSS_SHARE_MAX_LEN;
		TEE_MemMove(payloads[i] + payload_len++, &n, 1); // total #share, n
		TEE_MemMove(payloads[i] + payload_len++, &i, 1); // share index, i
		for (int j = 0; j < n; j++) { // checksums
			TEE_MemMove(payloads[i] + payload_len, checksums[j], CHECKSUM_LEN);
			payload_len += CHECKSUM_LEN;
		}
		meta_len = strnlen(cmd->payload.gen_backup_req_cmd.metadata, KEY_METADATA_MAX_LEN);
		TEE_MemMove(payloads[i] + payload_len,
				cmd->payload.gen_backup_req_cmd.metadata, meta_len); // metadata
		payload_len += meta_len;

		// verify buddy's cert chain and extract public key
		EVP_PKEY *buddy_evp_key;
		ret = verify_cert_chain(
				cmd->payload.gen_backup_req_cmd.buddies_cert_chains[i], false,
				&buddy_evp_key);
		if (ret != PEBBLE_STATUS_SUCCESS) {
			PEBBLE_LOG("%s: failed to verify cert chain", __FUNCTION__);
			goto exit;
		}
		ret = convert_evp_pkey_to_rsa_info(buddy_evp_key, &buddy_key, kid);
		if (ret != PEBBLE_STATUS_SUCCESS) {
			PEBBLE_LOG("%s: failed to convert_evp_pkey_to_rsa_info",
					__FUNCTION__);
			goto exit;
		}
		// wrap payload in JWE and JWS
		jwe_len = sizeof(jwe_out);
		jws_len = sizeof(jws_out);
		memset(jwe_out, 0, sizeof(jwe_out));
		memset(jws_out, 0, sizeof(jws_out));
		// encrypt with buddy i's public key (JWE)
		r = create_jwe_payload((uint8_t*) payloads[i], payload_len, &buddy_key,
				jwe_out, &jwe_len);
		if (r != JWE_OK) {
			PEBBLE_LOG("%s: create_jwe_payload failed: %d", __FUNCTION__, r);
			ret = PEBBLE_CREATE_JWE_FAILED;
			goto exit;
		}
//		printLongRange(jwe_out, jwe_len, "JWE");
		// sign with my private key (JWS)
		r = create_jws_payload((uint8_t*) jwe_out, jwe_len, &self_key, jws_out,
				&jws_len);
		if (r != JWS_OK) {
			PEBBLE_LOG("%s: create_jws_payload failed: %d", __FUNCTION__, r);
			ret = PEBBLE_CREATE_JWS_FAILED;
			goto exit;
		}
//		printLongRange(jws_out, jws_len, "JWS");

		TEE_MemMove(resp->payload.gen_backup_req_resp.backup_reqs_jws[i],
				jws_out, jws_len);
		TEE_MemMove(
				resp->payload.gen_backup_req_resp.backup_reqs_jws[i] + jws_len,
				"\0", 1);
	}
	resp->payload.gen_backup_req_resp.n = n;

	ret = PEBBLE_STATUS_SUCCESS;
	exit:
	memset(master_key, 0, sizeof(master_key));
	memset(shares, 0, sizeof(shares));
	memset(payloads, 0, sizeof(payloads));
	return ret;
}

/**
 * Input: backup request JWS string
 * Output: backup ack JWS string, wrapped key share, and some meta info
 */
pebble_return_code_t handle_backup_req_and_gen_backup_ack(tci_message_t *cmd,
		tci_message_t *resp) {
	pebble_return_code_t ret;
	if (!drk_ready) {
		ret = PEBBLE_DRK_NOT_READY;
		goto exit;
	}
	uint32_t jwt_len = JWE_JWS_MAX_SIZE;
	uint32_t payload_len = SHARE_PAYLOAD_LEN;
	uint8_t n, i, r, *pos, jwt[JWE_JWS_MAX_SIZE-1];
	uint8_t payload[SHARE_PAYLOAD_LEN];
	//uint8_t * share, * checksum;char * checksums;
	uint8_t share[SSS_SHARE_MAX_LEN], checksum[CHECKSUM_LEN],
			wshare[SK_WRAP_KEY_LEN];

	// verify buddy's cert chain and extract public key
	rsa_key_info_t buddy_key = { 0 };
	EVP_PKEY *buddy_evp_key;
	ret = verify_cert_chain(cmd->payload.gen_backup_ack_cmd.buddy_cert_chain,
			false, &buddy_evp_key);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG("%s: failed to verify buddy cert chain", __FUNCTION__);
		goto exit;
	}
	ret = convert_evp_pkey_to_rsa_info(buddy_evp_key, &buddy_key, kid);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG("%s: failed to convert_evp_pkey_to_rsa_info", __FUNCTION__);
		goto exit;
	}

	// extract payload from backup request JWS(JWE(payload, my key), buddy key)
	r = extract_jws_payload(
			(uint8_t*) cmd->payload.gen_backup_ack_cmd.backup_req_jws,
			strnlen(cmd->payload.gen_backup_ack_cmd.backup_req_jws, JWE_JWS_MAX_SIZE), &buddy_key,
			jwt, &jwt_len);
	if (r != JWS_OK) {
		PEBBLE_LOG("%s: extract_jws_payload failed: %d", __FUNCTION__, r);
		ret = PEBBLE_EXTRACT_JWS_FAILED;
		goto exit;
	}
//	printLongRange(jwt, jwt_len, "-JWS");
	r = extract_jwe_payload(jwt, jwt_len, &self_key, payload, &payload_len);
	if (r != JWE_OK) {
		PEBBLE_LOG("%s: extract_jwe_payload failed: %d", __FUNCTION__, r);
		ret = PEBBLE_EXTRACT_JWE_FAILED;
		goto exit;
	}
//	printLongRange(payload, payload_len, "-JWE"); // share and checksums are unsigned chars, metadata is printable chars

	/** parse payload based on our customized format:
	 *  "[share][share number n][share index i][all shares’ checksums concatenated][metadata_json_for_NWD_use]"
	 *     33 B      1 B             1 B            32*n B                             len doesn't matter
	 */
//	share = (uint8_t *) strtok((char *) payload, "|");
//	share_len = strnlen((const char*) share, SSS_SHARE_MAX_LEN);
//	i = (uint32_t) strtol(strtok(NULL, "|"), (char **)NULL, 10);
//	checksums = strtok(NULL, "|");
//	TEE_MemMove(checksum, checksums + i * CHECKSUM_LEN, CHECKSUM_LEN);
	TEE_MemMove(share, payload, SSS_SHARE_MAX_LEN); //get share
	pos = payload + SSS_SHARE_MAX_LEN;
	n = *(pos++); // get share index
	i = *(pos++); // get share index
	TEE_MemMove(checksum, pos + i * CHECKSUM_LEN, CHECKSUM_LEN);
	PEBBLE_LOG("share index i=%d", i);

	// verify share's checksum
	if (1 != verify_checksum(share, SSS_SHARE_MAX_LEN, checksum)) {
		PEBBLE_LOG("%s: verify_checksum failed", __FUNCTION__);
		printRangeHex(checksum, CHECKSUM_LEN, "wrong chksm");
		ret = PEBBLE_SHARE_CHECKSUM_MISMATCH;
		goto exit;
	}

	// wrap the share
	uint32_t wshare_len = SK_WRAP_KEY_LEN, checksums_len = n * CHECKSUM_LEN;
	uint32_t meta_len = payload_len - SSS_SHARE_MAX_LEN - 2 - checksums_len;
	ret = wrap(share, SSS_SHARE_MAX_LEN, wshare, &wshare_len);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG("%s: failed to wrap the share with TA", __FUNCTION__);
		ret = PEBBLE_WRAP_FAILED;  // need to add to tz_pebble_interface.h
		goto exit;
	}
	TEE_MemMove(resp->payload.gen_backup_ack_resp.wrapped_share.buf, wshare,
			wshare_len);
	resp->payload.gen_backup_ack_resp.wrapped_share.len = wshare_len;
	resp->payload.gen_backup_ack_resp.n = n;
	resp->payload.gen_backup_ack_resp.i = i;
	TEE_MemMove(resp->payload.gen_backup_ack_resp.checksums.buf, pos,
			checksums_len);
	resp->payload.gen_backup_ack_resp.checksums.len = checksums_len;
	TEE_MemMove(resp->payload.gen_backup_ack_resp.metadata, pos + checksums_len,
			meta_len);
	if (meta_len > KEY_METADATA_MAX_LEN - 1) {
		PEBBLE_LOG("%s: metadata is too long: %d bytes", __FUNCTION__, meta_len);
		ret = 1;
		goto exit;
	}
	TEE_MemMove(resp->payload.gen_backup_ack_resp.metadata + meta_len, "\0", 1);

	// generate the backup ack as a JWS string with "OK" as payload
	memset(jwt, 0, sizeof(jwt));
	jwt_len = sizeof(jwt);
	r = create_jws_payload((uint8_t*) "OK", 2, &self_key, jwt, &jwt_len);
	if (r != JWS_OK) {
		PEBBLE_LOG("%s: create_jws_payload failed: %d", __FUNCTION__, r);
		ret = PEBBLE_CREATE_JWS_FAILED;
		goto exit;
	}
//	printLongRange(jwt, jwt_len, "JWS");
	TEE_MemMove(resp->payload.gen_backup_ack_resp.backup_ack_jws, jwt, jwt_len);
	TEE_MemMove(resp->payload.gen_backup_ack_resp.backup_ack_jws + jwt_len,
			"\0", 1);

	ret = PEBBLE_STATUS_SUCCESS;
	exit:
	memset(share, 0, sizeof(share));
	memset(&buddy_key, 0, sizeof(buddy_key));
	memset(buddy_evp_key, 0, sizeof(EVP_PKEY));
	return ret;
}

/**
 * Input: key alias of the master key to be recovered
 * Output: recovery request JWS string
 */
pebble_return_code_t gen_recovery_req(tci_message_t *cmd, tci_message_t *resp) {
	pebble_return_code_t ret;
	if (!drk_ready) {
		ret = PEBBLE_DRK_NOT_READY;
		goto exit;
	}
	// create a JWS string: JWS(key alias, key)
	uint32_t r, jwt_len = JWE_JWS_MAX_SIZE-1;
	uint8_t *jwt =
			(uint8_t*) resp->payload.gen_recovery_req_resp.recovery_req_jws;
	r = create_jws_payload(
			(uint8_t*) cmd->payload.gen_recovery_req_cmd.key_alias,
			strnlen(cmd->payload.gen_recovery_req_cmd.key_alias, KEY_ALIAS_LEN), &self_key, jwt,
			&jwt_len);
	if (r != JWS_OK) {
		PEBBLE_LOG("%s: create_jws_payload failed: %d", __FUNCTION__, r);
		ret = 1; //PEBBLE_CREATE_JWS_FAILED;
		goto exit;
	}
//	printLongRange(jwt, jwt_len, "JWS");
	TEE_MemMove(jwt + jwt_len, "\0", 1);

	ret = PEBBLE_STATUS_SUCCESS;
	exit: return ret;
}

/**
 * Input: recovery request JWS string from buddy device, wrapped key share and some meta info, buddy's certificate chain
 * Output: recovery ack JWS string (containing a key share and some meta info), error message
 */
pebble_return_code_t handle_recovery_req_and_gen_recovery_ack(
		tci_message_t *cmd, tci_message_t *resp) {
	if (cmd->payload.gen_recovery_ack_cmd.wrapped_share.len > SK_WRAP_KEY_LEN ||
			cmd->payload.gen_recovery_ack_cmd.checksums.len > MAX_BUDDIES * CHECKSUM_LEN) {
		PEBBLE_LOG("%s: len from user input exceeds the limit", __FUNCTION__);
		return PEBBLE_INVALID_USER_INPUT_LEN;
	}
	if (cmd->payload.gen_recovery_ack_cmd.n > MAX_BUDDIES ||
			cmd->payload.gen_recovery_ack_cmd.i >= cmd->payload.gen_recovery_ack_cmd.n) {
		PEBBLE_LOG("%s: invalid user input. Ensure 0 <= i < n <= %d", __FUNCTION__, MAX_BUDDIES);
		return PEBBLE_INVALID_USER_INPUT_LEN;
	}
	pebble_return_code_t ret;
	if (!drk_ready) {
		ret = PEBBLE_DRK_NOT_READY;
		goto exit;
	}
	uint8_t r, alias[KEY_ALIAS_LEN], payload[SHARE_PAYLOAD_LEN];
	uint32_t alias_len = KEY_ALIAS_LEN, payload_len = SHARE_PAYLOAD_LEN;

	// verify buddy's cert chain and extract public key
	rsa_key_info_t buddy_key = { 0 };
	EVP_PKEY *buddy_evp_key;
	ret = verify_cert_chain(cmd->payload.gen_recovery_ack_cmd.buddy_cert_chain,
			false, &buddy_evp_key);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG("%s: failed to verify buddy cert chain", __FUNCTION__);
		goto exit;
	}
	ret = convert_evp_pkey_to_rsa_info(buddy_evp_key, &buddy_key, kid);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG("%s: failed to convert_evp_pkey_to_rsa_info", __FUNCTION__);
		goto exit;
	}

	// verify JWS signature
	r = extract_jws_payload(
			(uint8_t*) cmd->payload.gen_recovery_ack_cmd.recovery_req_jws,
			strnlen(cmd->payload.gen_recovery_ack_cmd.recovery_req_jws, JWE_JWS_MAX_SIZE),
			&buddy_key, alias, &alias_len);
	if (r != JWS_OK) {
		PEBBLE_LOG("%s: extract_jws_payload failed: %d", __FUNCTION__, r);
		ret = PEBBLE_EXTRACT_JWS_FAILED;
		goto exit;
	}
//	printLongRange(alias, alias_len, "-JWS");

	// unwrap the wrapped share
	uint32_t wshare_len = cmd->payload.gen_recovery_ack_cmd.wrapped_share.len;
	ret = unwrap(cmd->payload.gen_recovery_ack_cmd.wrapped_share.buf,
			&wshare_len, 1, (uint8_t*) payload, &payload_len); // put share at the beginning of payload
	if (ret != PEBBLE_STATUS_SUCCESS || payload_len != SSS_SHARE_MAX_LEN) {
		PEBBLE_LOG("%s: failed to unwrap the provided share", __FUNCTION__);
		goto exit;
	}

	// prepare payload in this format: [share][n][i][checksums][metadata]
	uint32_t checksums_len = cmd->payload.gen_recovery_ack_cmd.checksums.len;
	uint32_t meta_len = strnlen(cmd->payload.gen_recovery_ack_cmd.metadata, KEY_METADATA_MAX_LEN);
	TEE_MemMove(payload + payload_len++, &cmd->payload.gen_recovery_ack_cmd.n,
			1);
	TEE_MemMove(payload + payload_len++, &cmd->payload.gen_recovery_ack_cmd.i,
			1);
	TEE_MemMove(payload + payload_len,
			cmd->payload.gen_recovery_ack_cmd.checksums.buf, checksums_len);
	payload_len += checksums_len;
	TEE_MemMove(payload + payload_len,
			cmd->payload.gen_recovery_ack_cmd.metadata, meta_len);
	payload_len += meta_len;

	// encrypt payload with JWE and sign it with JWS:  JWS( JWE(payload, buddy_key), self_key)
	uint32_t jwe_len = JWE_JWS_MAX_SIZE, jws_len = JWE_JWS_MAX_SIZE-1;
	uint8_t jwe_out[JWE_JWS_MAX_SIZE], jws_out[JWE_JWS_MAX_SIZE];
	r = create_jwe_payload(payload, payload_len, &buddy_key, jwe_out, &jwe_len);
	if (r != JWE_OK) {
		PEBBLE_LOG("%s: create_jwe_payload failed: %d", __FUNCTION__, r);
		ret = PEBBLE_CREATE_JWE_FAILED;
		goto exit;
	}
//	printLongRange(jwe_out, jwe_len, "JWE");
	r = create_jws_payload(jwe_out, jwe_len, &self_key, jws_out, &jws_len);
	if (r != JWS_OK) {
		PEBBLE_LOG("%s: create_jws_payload failed: %d", __FUNCTION__, r);
		ret = PEBBLE_CREATE_JWS_FAILED;
		goto exit;
	}
//	printLongRange(jws_out, jws_len, "JWS");
	TEE_MemMove(resp->payload.gen_recovery_ack_resp.recovery_ack_jws, jws_out,
			jws_len);
	TEE_MemMove(resp->payload.gen_recovery_ack_resp.recovery_ack_jws + jws_len,
			"\0", 1);

	ret = PEBBLE_STATUS_SUCCESS;
	exit:
	memset(payload, 0, sizeof(payload));
	memset(&buddy_key, 0, sizeof(buddy_key));
	memset(buddy_evp_key, 0, sizeof(EVP_PKEY));
	return ret;
}

/**
 * Input: recovery ack JWS string (containing a key share wrapped by JWE), buddy's certificate chain
 * Output: key share wrapped by self TA, error message
 */
pebble_return_code_t handle_recovery_ack(tci_message_t *cmd,
		tci_message_t *resp) {
	pebble_return_code_t ret;
	if (!drk_ready) {
		ret = PEBBLE_DRK_NOT_READY;
		goto exit;
	}
	uint32_t jwt_len = JWE_JWS_MAX_SIZE;
	uint32_t payload_len = SHARE_PAYLOAD_LEN;
	uint8_t n, i, r, *pos, jwt[JWE_JWS_MAX_SIZE];
	uint8_t payload[SHARE_PAYLOAD_LEN];
	uint8_t share[SSS_SHARE_MAX_LEN], checksum[CHECKSUM_LEN],
			wshare[SK_WRAP_KEY_LEN];

	// verify buddy's cert chain and extract public key
	rsa_key_info_t buddy_key = { 0 };
	EVP_PKEY *buddy_evp_key;
	ret = verify_cert_chain(
			cmd->payload.handle_recovery_ack_cmd.buddy_cert_chain, false,
			&buddy_evp_key);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG("%s: failed to verify buddy cert chain", __FUNCTION__);
		goto exit;
	}
	ret = convert_evp_pkey_to_rsa_info(buddy_evp_key, &buddy_key, kid);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG("%s: failed to convert_evp_pkey_to_rsa_info", __FUNCTION__);
		goto exit;
	}

	// extract payload from backup request JWS(JWE(payload, my key), buddy key)
	r = extract_jws_payload(
			(uint8_t*) cmd->payload.handle_recovery_ack_cmd.recovery_ack_jws,
			strnlen(cmd->payload.handle_recovery_ack_cmd.recovery_ack_jws, JWE_JWS_MAX_SIZE),
			&buddy_key, jwt, &jwt_len);
	if (r != JWS_OK) {
		PEBBLE_LOG("%s: extract_jws_payload failed: %d", __FUNCTION__, r);
		ret = PEBBLE_EXTRACT_JWS_FAILED;
		goto exit;
	}
//	printLongRange(jwt, jwt_len, "-JWS");
	r = extract_jwe_payload(jwt, jwt_len, &self_key, payload, &payload_len);
	if (r != JWE_OK) {
		PEBBLE_LOG("%s: extract_jwe_payload failed: %d", __FUNCTION__, r);
		ret = PEBBLE_EXTRACT_JWE_FAILED;
		goto exit;
	}
//	printLongRange(payload, payload_len, "-JWE"); // share and checksums are unsigned chars, metadata is printable chars

	/** parse payload based on our customized format:
	 *  "[share][share index][all shares’ checksums concatenated][metadata_json_for_NWD_use]"
	 *     33 B      1 B         32*n B                             len doesn't matter
	 */
	TEE_MemMove(share, payload, SSS_SHARE_MAX_LEN); //get share
	pos = payload + SSS_SHARE_MAX_LEN;
	n = *(pos++); // get share index
	i = *(pos++); // get share index
	TEE_MemMove(checksum, pos + i * CHECKSUM_LEN, CHECKSUM_LEN);
	PEBBLE_LOG("share index i=%d", i);

	// verify share's checksum
	if (1 != verify_checksum(share, SSS_SHARE_MAX_LEN, checksum)) {
		PEBBLE_LOG("%s: verify_checksum failed", __FUNCTION__);
		printRangeHex(checksum, CHECKSUM_LEN, "wrong chksm");
		ret = PEBBLE_SHARE_CHECKSUM_MISMATCH;
		goto exit;
	}

	// wrap the share
	uint32_t wshare_len = SK_WRAP_KEY_LEN, checksums_len = n * CHECKSUM_LEN;
	uint32_t meta_len = payload_len - SSS_SHARE_MAX_LEN - 2 - checksums_len;
	ret = wrap(share, SSS_SHARE_MAX_LEN, wshare, &wshare_len);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG("%s: failed to wrap the share with TA", __FUNCTION__);
		ret = PEBBLE_WRAP_FAILED;  // need to add to tz_pebble_interface.h
		goto exit;
	}
	TEE_MemMove(resp->payload.handle_recovery_ack_resp.wrapped_share.buf,
			wshare, wshare_len);
	resp->payload.handle_recovery_ack_resp.wrapped_share.len = wshare_len;
	resp->payload.handle_recovery_ack_resp.n = n;
	resp->payload.handle_recovery_ack_resp.i = i;
	TEE_MemMove(resp->payload.handle_recovery_ack_resp.checksums.buf, pos,
			checksums_len);
	resp->payload.handle_recovery_ack_resp.checksums.len = checksums_len;
	TEE_MemMove(resp->payload.handle_recovery_ack_resp.metadata,
			pos + checksums_len, meta_len);
	if (meta_len > KEY_METADATA_MAX_LEN - 1) {
		PEBBLE_LOG("%s: metadata is too long: %d bytes", __FUNCTION__, meta_len);
		ret = 1;
		goto exit;
	}
	TEE_MemMove(resp->payload.handle_recovery_ack_resp.metadata + meta_len,
			"\0", 1);

	ret = PEBBLE_STATUS_SUCCESS;
	exit:
	memset(share, 0, sizeof(share));
	memset(payload, 0, sizeof(payload));
	memset(&buddy_key, 0, sizeof(buddy_key));
	memset(buddy_evp_key, 0, sizeof(EVP_PKEY));
	return ret;
}

/**
 * reconstruct a master key, and return one copy wrapped by TA, another copy wrapped by Android KeyStore public key
 * Input: all shares for the same master key, AES-GCM AAD(gcm_aad), Android Key Store certificate chain
 * Output: master key wrapped by TA, master key encrypted by AES-GCM (secret_ct, gcm_iv, gcm_tag), aes key encrypted (wrapped_gcm_key) by provided public key
 */
pebble_return_code_t reconstruct_key(tci_message_t *cmd, tci_message_t *resp) {
	if (cmd->payload.reconstruct_key_cmd.wshare_len > SK_WRAP_KEY_LEN
			|| cmd->payload.reconstruct_key_cmd.master_key_req.gcm_aad.len > WRAPPING_AES_AAD_LEN
			|| cmd->payload.reconstruct_key_cmd.k > MAX_BUDDIES) {
		PEBBLE_LOG("%s: len from user input exceeds the limit", __FUNCTION__);
		return PEBBLE_INVALID_USER_INPUT_LEN;
	}
	pebble_return_code_t ret;
	uint8_t k = cmd->payload.reconstruct_key_cmd.k;
	uint32_t wshare_len = cmd->payload.reconstruct_key_cmd.wshare_len,
			share_len, mk_len;
	uint8_t shares[MAX_BUDDIES][SSS_SHARE_MAX_LEN], master_key[SSS_KEY_MAX_LEN],
			share_tmp[SK_WRAP_KEY_LEN];
	EVP_PKEY *pkey = NULL;
	uint8_t *key_wrapped_ta;
	uint32_t key_wrapped_ta_len;
	uint8_t *secret_ct;
	uint32_t secret_ct_len;
	uint8_t *wrapped_gcm_key;
	uint32_t wrapped_gcm_key_len;

	// Validate x509 certificates chains
	PEBBLE_LOG_DEBUG("Validate cert chain");
#ifdef LOCAL_TEST
	const char *test_nonce = "hello world";
	ret = verify_aks_cert_chain(
			cmd->payload.reconstruct_key_cmd.master_key_req.cert_chain, &pkey,
			(uint8_t*) test_nonce, strnlen(test_nonce, ATN_NONCE_MAX_LEN));
#else
    ret = verify_aks_cert_chain(cmd->payload.gen_master_key_cmd.cert_chain, &pkey, nonce, nonce_len);
#endif
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG("verify_aks_cert_chain FAIL");
		PEBBLE_LOG_DEBUG("verify_aks_cert_chain FAIL ret = %d", ret);
		return PEBBLE_WRAP_FAILED;
	}

	// unwrap each share
	for (int i = 0; i < k; i++) {
		share_len = SK_WRAP_KEY_LEN;
		ret = unwrap(cmd->payload.reconstruct_key_cmd.wrapped_shares[i],
				&wshare_len, 1, share_tmp, &share_len); // put share at the beginning of payload
		if (ret != PEBBLE_STATUS_SUCCESS || share_len != SSS_SHARE_MAX_LEN) {
			PEBBLE_LOG("%s: failed to unwrap the provided share %d",
					__FUNCTION__, i);
			goto exit;
		}
		TEE_MemMove(shares[i], share_tmp, SSS_SHARE_MAX_LEN);
	}

	// combine the shares into a master key
	join(shares, share_len, k, master_key, &mk_len);
	if (mk_len != SSS_KEY_MAX_LEN) {
		PEBBLE_LOG("%s: reconstructed key len is %d, not %d", __FUNCTION__,
				mk_len, SSS_KEY_MAX_LEN);
		ret = PEBBLE_SSS_JOIN_ERROR;
		goto exit;
	}
//	printLongRangeHex(master_key, SK_MASTER_KEY_LEN, "master_key before wrap");

	// wrap the master key into 2 copies: one is wrapped by TA, the other is wrapped by Android KeyStore
	key_wrapped_ta = TEE_Malloc(SK_WRAP_KEY_LEN, 0);
	secret_ct = TEE_Malloc(SK_WRAP_KEY_LEN, 0);
	wrapped_gcm_key = TEE_Malloc(WRAPPING_RSA_N_BYTES_LEN, 0);
	if (key_wrapped_ta == NULL || secret_ct == NULL || wrapped_gcm_key == NULL) {
		ret = PEBBLE_ALLOC_ERROR;
		PEBBLE_LOG("Error to alloc memory");
		goto exit;
	}
	key_wrapped_ta_len = SK_WRAP_KEY_LEN;
	secret_ct_len = SK_WRAP_KEY_LEN;
	wrapped_gcm_key_len = WRAPPING_RSA_N_BYTES_LEN;

	// 1. wrap it with TA
	ret = wrap((uint8_t*) master_key, sizeof(master_key), key_wrapped_ta,
			&key_wrapped_ta_len);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG(
				"generate_master_key_and_wrap_it - failed to wrap the master key with TA");
		ret = PEBBLE_WRAP_FAILED;  // need to add to tz_pebble_interface.h
		goto exit;
	}

	// 2. wrap it with Android KeyStore public key
	ret = two_step_wrap_with_android_keystore_public_key(master_key,
			sizeof(master_key), pkey,
			cmd->payload.reconstruct_key_cmd.master_key_req.gcm_aad.buf,
			cmd->payload.reconstruct_key_cmd.master_key_req.gcm_aad.len,
			secret_ct, &secret_ct_len,
			resp->payload.reconstruct_key_resp.master_key_resp.gcm_iv,
			WRAPPING_AES_IV_LEN,
			resp->payload.reconstruct_key_resp.master_key_resp.gcm_tag,
			WRAPPING_AES_TAG_LEN, wrapped_gcm_key, &wrapped_gcm_key_len);
	if (ret != PEBBLE_STATUS_SUCCESS) {
		PEBBLE_LOG(
				"generate_master_key_and_wrap_it - failed to wrap the master key with Android KeyStore public key");
		ret = PEBBLE_WRAP_FAILED;  // need to add to tz_pebble_interface.h
		goto exit;
	}
//	printLongRangeHex(master_key, sizeof(master_key), "master_key");
//	printLongRangeHex(key_wrapped_ta, key_wrapped_ta_len, "key_wrapped_ta");
//	printLongRangeHex(secret_ct, secret_ct_len, "secret_ct");
//	printLongRangeHex(wrapped_gcm_key, wrapped_gcm_key_len, "wrapped_gcm_key");

	resp->payload.reconstruct_key_resp.master_key_resp.master_key_wrapped_ta.len =
			key_wrapped_ta_len;
	TEE_MemMove(
			resp->payload.reconstruct_key_resp.master_key_resp.master_key_wrapped_ta.buf,
			key_wrapped_ta, key_wrapped_ta_len);
	resp->payload.reconstruct_key_resp.master_key_resp.secret_ct.len =
			secret_ct_len;
	TEE_MemMove(
			resp->payload.reconstruct_key_resp.master_key_resp.secret_ct.buf,
			secret_ct, secret_ct_len);
	resp->payload.reconstruct_key_resp.master_key_resp.wrapped_gcm_key.len =
			wrapped_gcm_key_len;
	TEE_MemMove(
			resp->payload.reconstruct_key_resp.master_key_resp.wrapped_gcm_key.buf,
			wrapped_gcm_key, wrapped_gcm_key_len);

	ret = PEBBLE_STATUS_SUCCESS;
	exit: if (key_wrapped_ta != NULL) {
		TEE_Free(key_wrapped_ta);
	}
	if (secret_ct != NULL) {
		TEE_Free(secret_ct);
	}
	if (wrapped_gcm_key != NULL) {
		TEE_Free(wrapped_gcm_key);
	}
	memset(shares, 0, sizeof(shares));
	memset(share_tmp, 0, sizeof(share_tmp));
	memset(master_key, 0, sizeof(master_key));
	return ret;
}