#include "qsee_log.h"
#include "qsee_stor.h"
#include "qsee_fuse.h"
#include "qsee_fs.h"
#include "qsee_core.h"

#include <em_ta.h>

static qsee_stor_device_handle_t gDevHandle;
static qsee_stor_client_handle_t gClientHandle;

static int em_qsee_init_rpmb()
{
	int ret = qsee_stor_device_init(QSEE_STOR_EMMC_RPMB, 0, &gDevHandle);
	if (ret != QSEE_STOR_SUCCESS) {
		LOGE("%s: Failed to init RPMB device(%d)\n", __func__, ret);
		return EM_ERR_EM_QSEE_RPMB_INIT;
	}

	ret = qsee_stor_open_partition(&gDevHandle, EM_RPMB_QSEE_PARTITION_ID, &gClientHandle);
	if (ret == QSEE_STOR_PARTI_NOT_FOUND_ERROR) {
		LOGI("%s: There is no partition engmode, add it.\n", __func__);

		qsee_stor_device_info_t device_info;
		uint16_t num_rpmb_sectors = EM_RPMB_RESERVED_SECTOR_NUM;

		ret = qsee_stor_device_get_info(&gDevHandle, &device_info);
		if (ret != QSEE_STOR_SUCCESS) {
			LOGE("%s: Failed to get device info(%d)\n", __func__, ret);
			num_rpmb_sectors = EM_RPMB_RESERVED_SECTOR_NUM;
		} else {
			LOGI("%s: Total RPMB sectors (%u), Available RPMB sectors (%u)\n", __func__,
				 device_info.total_sectors, device_info.available_sectors);
			num_rpmb_sectors = (device_info.available_sectors > EM_RPMB_RESERVED_SECTOR_NUM)
								   ? EM_RPMB_RESERVED_SECTOR_NUM
								   : device_info.available_sectors;
		}

		ret = qsee_stor_add_partition(&gDevHandle, EM_RPMB_QSEE_PARTITION_ID, num_rpmb_sectors);
		if (ret != QSEE_STOR_SUCCESS) {
			LOGE("%s: Failed to make new partition(%d)\n", __func__, ret);
			return EM_ERR_EM_QSEE_RPMB_ADD_PARTITION;
		}

		ret = qsee_stor_open_partition(&gDevHandle, EM_RPMB_QSEE_PARTITION_ID, &gClientHandle);
		if (ret != QSEE_STOR_SUCCESS) {
			LOGE("%s: Failed to open partition again(%d)\n", __func__, ret);
			return EM_ERR_EM_QSEE_RPMB_OPEN_PART;
		}

	} else if (ret != QSEE_STOR_SUCCESS) {
		LOGE("%s: Failed to open partition(%d)\n", __func__, ret);
		return EM_ERR_EM_QSEE_RPMB_NOT_OPENED;
	}

	return EM_SUCCESS;
}

static int em_qsee_rpmb_read(uint32_t sector_pos, uint8_t *data) {
	int ret;

	ret = qsee_stor_read_sectors(&gClientHandle, sector_pos, 1, data);
	if (ret != QSEE_STOR_SUCCESS) {
		LOGE("%s: Failed to read data from rpmb(0x%08x) - try again\n", __func__, ret);
		ret = qsee_stor_read_sectors(&gClientHandle, sector_pos, 1, data);
		if (ret != QSEE_STOR_SUCCESS) {
			LOGE("%s: Failed to read data from rpmb(0x%08x) - twice\n", __func__, ret);
			ret = EM_ERR_EM_QSEE_RPMB_READ_SECTOR;
			goto out;
		}
	}

	ret = EM_SUCCESS;

out:
	return ret;
}

static int em_qsee_rpmb_write(uint32_t sector_pos, uint8_t *data) {
	int ret;

	ret = qsee_stor_write_sectors(&gClientHandle, sector_pos, 1, data);
	if (ret != QSEE_STOR_SUCCESS) {
		LOGE("%s: Failed to write data to rpmb(%d) - try again\n", __func__, ret);
		ret = qsee_stor_write_sectors(&gClientHandle, sector_pos, 1, data);
		if (ret != QSEE_STOR_SUCCESS) {
			LOGE("%s: Failed to write data to rpmb(%d) twice\n", __func__, ret);
			ret = EM_ERR_EM_QSEE_RPMB_WRITE_SECTOR;
			goto out;
		}
	}

	ret = qsee_stor_write_sectors(&gClientHandle, sector_pos + 1, 1, data);
	if (ret != QSEE_STOR_SUCCESS) {
		LOGE("%s: Failed to write data to rpmb(%d) - try again\n", __func__, ret);
		ret = qsee_stor_write_sectors(&gClientHandle, sector_pos + 1, 1, data);
		if (ret != QSEE_STOR_SUCCESS) {
			LOGE("%s: Failed to write data to rpmb(%d) twice\n", __func__, ret);
			ret = EM_ERR_EM_QSEE_RPMB_WRITE_SECTOR_BACKUP;
			goto out;
		}
	}

	ret = EM_SUCCESS;

out:
	return ret;
}

int em_read_core(uint8_t *buf, uint32_t buf_len) {
	int ret;
	uint8_t core[EM_RPMB_SINGLE_BLOCK_SIZE*2] = {};
	uint8_t key[EM_LEN_KEY_CORE_V20] = {};
	uint8_t iv[EM_LEN_IV] = {};
	uint32_t plaintext_len = 0;
	uint32_t encrypt_len = 0;
	uint32_t unencrypt_len = 0;

	// Param check
	EM_CHECK_NULL(__func__, EM_ERR_EM_READ_CORE, buf);

	if (buf_len < EM_RPMB_SINGLE_BLOCK_SIZE) {
		LOGE("%s : buf_len is not enough(%u)\n", __func__, buf_len);
		ret = EM_ERR_EM_READ_CORE_QSEE_INVALID_ARG;
		goto out;
	}

	LOGI("%s : init rpmb\n", __func__);
	ret = em_qsee_init_rpmb();
	if (ret != EM_SUCCESS) {
		LOGE("%s: Failed to init RPMB(0x%08x)\n", __func__, ret);
		goto out;
	}

	// Read from RPMB
	LOGI("%s : read rpmb\n", __func__);
	ret = em_qsee_rpmb_read(0, core);
	if (ret != EM_SUCCESS) {
		LOGI("%s : Failed to read core, try again(0x%08x)\n", __func__, ret);
		ret = em_qsee_rpmb_read(1, core);
		if (ret != EM_SUCCESS) {
			LOGE("%s : Failed to read again(0x%08x)\n", __func__, ret);
			goto out;
		}
	}

	if (em_is_all_zero((uint8_t *)core, EM_RPMB_SINGLE_BLOCK_SIZE*2) == EM_SUCCESS) {
		LOGE("%s: The block is all zero\n", __func__);
		ret = EM_ERR_EM_READ_CORE_ALL_ZERO;
		goto out;
	}

	// Get Key
	LOGI("%s : kdf\n", __func__);
	ret = em_crypto_kdf(key, EM_LEN_KEY_CORE_V20, iv, EM_LEN_IV);
	if (ret != EM_SUCCESS) {
		LOGI("%s : Failed to derive key and iv(0x%08x)\n", __func__, ret);
		goto out;
	}

	memcpy(buf, core, EM_RPMB_SINGLE_BLOCK_SIZE);
	unencrypt_len = EM_LEN_MAGIC_CORE_V20 + EM_LEN_GCM_TAG_CORE_V20;
	encrypt_len = EM_RPMB_SINGLE_BLOCK_SIZE - unencrypt_len;

	// Unwap data
	LOGI("%s : decrypt\n", __func__);
	ret = em_crypto_aes_256_gcm_decrypt(core + unencrypt_len, encrypt_len, buf + unencrypt_len, &plaintext_len, key,
					    EM_LEN_KEY_CORE_V20, iv, EM_LEN_IV, core + EM_LEN_MAGIC_CORE_V20,
					    EM_LEN_GCM_TAG_CORE_V20);
	if (ret != EM_SUCCESS) {
		LOGE("%s : Failed to decrypt core(0x%08x)\n", __func__, ret);
		goto out;
	}

	if (plaintext_len != encrypt_len) {
		LOGE("%s : Core is decrypted, but the output length is unexpected(%u/%u)",
			 __func__, plaintext_len, encrypt_len);
		ret = EM_ERR_EM_READ_CORE_QSEE_MISMATCH_LEN;
		goto out;
	}

	if (em_is_all_zero((uint8_t *)buf + unencrypt_len, EM_RPMB_SINGLE_BLOCK_SIZE - unencrypt_len) == EM_SUCCESS) {
		LOGE("%s: The block is all zero\n", __func__);
		ret = EM_ERR_EM_READ_CORE_ALL_ZERO;
		goto out;
	}
out:
	memset(key, 0, EM_LEN_KEY_CORE_V20);
	memset(iv, 0, EM_LEN_IV);

	return ret;
}

int em_write_core(uint8_t *buf, uint32_t buf_len) {
	int ret;
	uint8_t core[EM_RPMB_SINGLE_BLOCK_SIZE*2] = {};
	uint8_t key[EM_LEN_KEY_CORE_V20] = {};
	uint8_t iv[EM_LEN_IV] = {};
	uint32_t ciphertext_len = 0;
	uint32_t encrypt_len = 0;
	uint32_t unencrypt_len = 0;

	// Param check
	EM_CHECK_NULL(__func__, EM_ERR_EM_WRITE_CORE, buf);

	if (buf_len < EM_RPMB_SINGLE_BLOCK_SIZE) {
		LOGE("%s : buf_len is not enough(%u)\n", __func__, buf_len);
		ret = EM_ERR_EM_WRITE_CORE_QSEE_INVALID_ARG;
		goto out;
	}

	LOGI("%s: em_qsee_init rpmb\n", __func__);
	ret = em_qsee_init_rpmb();
	if (ret != EM_SUCCESS) {
		LOGE("%s: Failed to init RPMB(0x%08x)\n", __func__, ret);
		goto out;
	}

	// Get Key
	LOGI("%s : kdf\n", __func__);
	ret = em_crypto_kdf(key, EM_LEN_KEY_CORE_V20, iv, EM_LEN_IV);
	if (ret != EM_SUCCESS) {
		LOGI("%s : Failed to derive key and iv(0x%08x)\n", __func__, ret);
		goto out;
	}

	memcpy(core, buf, EM_RPMB_SINGLE_BLOCK_SIZE);
	unencrypt_len = EM_LEN_MAGIC_CORE_V20 + EM_LEN_GCM_TAG_CORE_V20;
	encrypt_len = EM_RPMB_SINGLE_BLOCK_SIZE - unencrypt_len;

	// Wrap data
	LOGI("%s : encrypt\n", __func__);
	ret = em_crypto_aes_256_gcm_encrypt(buf + unencrypt_len, encrypt_len,
										core + unencrypt_len, &ciphertext_len,
										key, EM_LEN_KEY_CORE_V20,
										iv, EM_LEN_IV,
										core + EM_LEN_MAGIC_CORE_V20, EM_LEN_GCM_TAG_CORE_V20);
	if (ret != EM_SUCCESS) {
		LOGE("%s : Failed to encrypt core(0x%08x)\n", __func__, ret);
		goto out;
	}

	if (ciphertext_len != encrypt_len) {
		LOGE("%s : Core is encrypted, but the output length is unexpected(%u/%u)",
			 __func__, ciphertext_len, encrypt_len);
		ret = EM_ERR_EM_WRITE_CORE_QSEE_MISMATCH_LEN;
		goto out;
	}

	// Read from RPMB
	LOGI("%s : write\n", __func__);
	ret = em_qsee_rpmb_write(0, core);
	if (ret != EM_SUCCESS) {
		LOGI("%s : Failed to write core(0x%08x)\n", __func__, ret);
		goto out;
	}

out:

	// TODO Jaesung
	// Zero write for core data?
	memset(key, 0, EM_LEN_KEY_CORE_V20);
	memset(iv, 0, EM_LEN_IV);

    return ret;
}

int em_check_provision()
{
	int ret;
	qsee_secctrl_secure_status_t status;

	LOGW("Checking provision start...\n");

	memset(&status, 0, sizeof(status));
	ret = qsee_get_secure_state(&status);
	if (ret != EM_SUCCESS) {
		LOGE("%s : Failed to get secure_state(%08x)\n", __func__, ret);
		ret = EM_ERR_EM_QSEE_RPMB_SECURE_STATE;
		goto out;
	}

	if ((status.value[0] & (1 << 5)) != 0) {
		LOGI("%s : RPMB key isn't provisioned(%08x, %08x, %08x)\n", __func__, ret, status.value[0],
		     status.value[1]);
		ret = EM_ERR_EM_CHECK_PROVISION_NOT_PROVISION;
		goto out;
	}

	LOGI("%s : %08x, %08x\n", __func__, status.value[0], status.value[1]);
	ret = EM_SUCCESS;
out:
	LOGW("Checking provision done...(0x%08x)\n", ret);
	return ret;
}