#include <string.h>

#include "vk_prepare_service.h"
#include "vk_error.h"
#include "vk_log.h"
#include "vk_data_struct.h"
#include "vk_interface.h"
#include "vk_table.h"
#include "vk_constants.h"
#include "vk_utils.h"
#include "crypto/vk_crypto.h"
#include "openssl/mem.h"

/*
	Vault Structure for VaultKeeper Service

	- Unsheltered vault data(32B) : system property(8B) + build version(24B)
	- Unsheltered vault reserved(88B) : VEK hash(32B) + AP S/N(12B) + VK ID(16B)
	- Sheltered vault(160B)
*/

int vk_prepare_service(target_t* target, cmd_req_t* req, cmd_rsp_t* rsp)
{
	int ret = VK_ERR_GENERAL;
	unsvault_t* vault = NULL;
	uint32_t rpmb_key_prov = 0;
	uint32_t custom_kernel = OEM_FLAG_CUSTOM_KERNEL;
	int i;

	char* it = NULL;
	char* delim = "@@@";
	char latest_bin_ver[SHORT_BUILD_VER_LEN] = {0,};
	char build_ver_history[MAX_BUILD_VER_FIELD_LEN] = {0,};
	uint8_t vault_eKey[VEK_LEN] = {0,};
	uint8_t hash_eKey[VEK_LEN] = {0,};
	uint8_t vk_id[VK_ID_LEN + 1] = {0,};
	char prop_value;

	if (target == NULL) {
		LOGE("%s: target is null\n", __func__);
		return VK_ERR_INVALID_ARGUMENT;
	}

	if (req == NULL) {
		LOGE("%s: req is null\n", __func__);
		return VK_ERR_INVALID_ARGUMENT;
	}

	if (rsp == NULL) {
		LOGE("%s: rsp is null\n", __func__);
		return VK_ERR_INVALID_ARGUMENT;
	}

	if (isAllZero(req->unsheltered, MAX_UNSHELTERED_DATA_LEN)) {
		LOGE("%s: Request unsheltered data is empty\n", __func__);
		ret = VK_ERR_INVALID_ARGUMENT;
		goto out;
	}

	if (req->unsheltered_len > MAX_UNSHELTERED_DATA_LEN) {
		LOGE("%s: Request unsheltered data size is invalid(%d)\n", __func__, req->unsheltered_len);
		ret = VK_ERR_INVALID_ARGUMENT;
		goto out;
	}

	ret = target->platform.check_rpmb_key_provisioning(&rpmb_key_prov);
	if (ret != VK_SUCCESS) {
		LOGE("%s: Failed to check rpmb key provisioning(%d)\n", __func__, ret);
		goto out;
	}

	// Write 'STEADY' partition when RPMB is not available.
	if (rpmb_key_prov == RPMB_NOT_PROVISIONED) {
		req->vault_level = VAULT_LEVEL1;
	} else {
		req->vault_level = VTAB[req->vtab_index].vault_level;
	}

	ret = target->platform.memory_alloc(&vault, sizeof(unsvault_t));
	if (ret != VK_SUCCESS || vault == NULL) {
		LOGE("%s: Failed memory alloc(%d)\n", __func__, ret);
		goto out;
	}

	ret = target->vault.read_unsheltered_vault(target, req, vault);
	if (ret != VK_SUCCESS) {
		LOGE("%s: Failed read vk vault(%d)\n", __func__, ret);
		goto out;
	}

	// 1. Save system property
	// PROP_TYPE_FIRST_API_LEVEL : Once written, It doesn't change
	prop_value = vault->unsheltered_data[PROP_TYPE_FIRST_API_LEVEL];
	if (prop_value != 'F' && prop_value != 'T' && prop_value != 0x0) {
		req->unsheltered[PROP_TYPE_FIRST_API_LEVEL] = prop_value;
	}

	memcpy(vault->unsheltered_data, req->unsheltered, MAX_PROP_FIELD_LEN);
	// <--- End 1. Save system property

	// 2. Check whether binary is custom kernel or not
	ret = target->platform.check_custom_kernel(&custom_kernel);
	if (ret != VK_SUCCESS) {
		LOGE("%s: Failed to check oem_sw_fuse(%d)\n", __func__, ret);
		goto out;
	}

	if (custom_kernel == OEM_FLAG_CUSTOM_KERNEL) {
		vault->unsheltered_data[PROP_TYPE_CUSTOM_KERNEL] = 'T';
	} else {
		vault->unsheltered_data[PROP_TYPE_CUSTOM_KERNEL] = 'F';
	}
	// <--- End 2. Check whether binary is custom kernel or not

	// 3. Update build binary version
	// version prefix - U:USER/F:FACTORY
	memcpy(build_ver_history, vault->unsheltered_data + MAX_PROP_FIELD_LEN, MAX_BUILD_VER_FIELD_LEN);

	// Copy from STEADY when downloading user binary first
	it = strstr(build_ver_history, delim);
	if (it == NULL && req->vault_level == VAULT_LEVEL2 && !isAllZero(req->preload + MAX_PROP_FIELD_LEN, MAX_BUILD_VER_FIELD_LEN)) {
		memcpy(build_ver_history, req->preload + MAX_PROP_FIELD_LEN, SHORT_BUILD_VER_LEN);
	}

	it = strstr(build_ver_history, delim);
	if (it == NULL) {
		// Update first.
		// ex) "" => 'UTB' => "UTB@@@"
		memcpy(build_ver_history, req->unsheltered + MAX_PROP_FIELD_LEN, SHORT_BUILD_VER_LEN);
		memcpy(build_ver_history + SHORT_BUILD_VER_LEN, delim, SHORT_BUILD_VER_LEN);
	} else {
		if (it == build_ver_history) {
			memcpy(latest_bin_ver, build_ver_history + MAX_BUILD_VER_FIELD_LEN - SHORT_BUILD_VER_LEN, SHORT_BUILD_VER_LEN);
		} else {
			memcpy(latest_bin_ver, it - SHORT_BUILD_VER_LEN, SHORT_BUILD_VER_LEN);
		}

		// Update only if it's not same version
		if(CRYPTO_memcmp(latest_bin_ver, req->unsheltered + MAX_PROP_FIELD_LEN, SHORT_BUILD_VER_LEN)) {
			// ex)
			// "UT1UT2@@@" => 'UTB' => "UT1UT2UTB@@@"
			// "UT1UT2UT5UT9@@@US1US2US3" => 'UTB' => "UT1UT2UT5UT9UTB@@@US2US3"
			// "UT1UT2UT5UT6UT7UT8UT9@@@" => 'UTB' => "@@@UT2UT5UT6UT7UT8UT9UTB"
 			// "UT1UT2UT5UT6UT7UT8UT9@@@" => 'UT9' => "UT1UT2UT5UT6UT7UT8UT9@@@"
			memcpy(it, req->unsheltered + MAX_PROP_FIELD_LEN, SHORT_BUILD_VER_LEN);

			if (it == (build_ver_history + MAX_BUILD_VER_FIELD_LEN - SHORT_BUILD_VER_LEN)) {
				it = build_ver_history;
			} else {
				it += SHORT_BUILD_VER_LEN;
			}

			memcpy(it, delim, SHORT_BUILD_VER_LEN);
		}
	}
	memcpy(vault->unsheltered_data + MAX_PROP_FIELD_LEN, build_ver_history, MAX_BUILD_VER_FIELD_LEN);
	// <--- End 3. Update build binary version

	// 4. Check & Save vault encryption key
	ret = VK_GET_SECURE_ITEM(VK_SECURE_ITEM_KEY, vault_eKey, VEK_LEN);
	if (ret != VK_SUCCESS) {
		LOGE("%s: Failed to get secure item(%d/%d)\n", __func__, VK_SECURE_ITEM_KEY, ret);
		goto out;
	}

	ret = vk_crypto_sha256(vault_eKey, sizeof(vault_eKey), hash_eKey);
	if (ret != VK_SUCCESS) {
		LOGE("%s: Failed to make a hash(%d)\n", __func__, ret);
		goto out;
	}

	vault->unsheltered_data[PROP_TYPE_VEK_DIFF] = 'F';

	if (isAllZero(vault->reserved, VEK_LEN)) {
		memcpy(vault->reserved, hash_eKey, VEK_LEN);
	} else {
		if (CRYPTO_memcmp(hash_eKey, vault->reserved, VEK_LEN) != 0) {
			LOGE("%s: Critical error! Mismatch vault eKey\n", __func__);
			vault->unsheltered_data[PROP_TYPE_VEK_DIFF] = 'T';
		}
	}
	// <--- 4. Check & Save vault encryption key

	// 5. Check & Save AP S/N
	// Ref. main.cpp :: preparing_vaultkeeper_service()
	vault->unsheltered_data[PROP_TYPE_AP_SN] = 'F';

	if (isAllZero(vault->reserved + VEK_LEN, AP_SN_LEN)) {
		memcpy(vault->reserved + VEK_LEN, req->unsheltered + MAX_PROP_FIELD_LEN + SHORT_BUILD_VER_LEN, AP_SN_LEN);
	} else {
		if (CRYPTO_memcmp(vault->reserved + VEK_LEN, req->unsheltered + MAX_PROP_FIELD_LEN + SHORT_BUILD_VER_LEN, AP_SN_LEN) != 0) {
			LOGE("%s: Mismatch AP S/N.\n", __func__);
			vault->unsheltered_data[PROP_TYPE_AP_SN] = 'T';
		}
	}
	// <----5. Check & Save AP S/N

	// 6. Check & Generate VK ID
	if (isAllZero(vault->reserved + VEK_LEN + AP_SN_LEN, VK_ID_LEN)) {
		ret = VK_GET_SECURE_ITEM(VK_SECURE_ITEM_NONCE, vk_id, VK_ID_LEN);
		if (ret != VK_SUCCESS) {
			LOGE("%s: Failed to get secure item(%d/%d)\n", __func__, VK_SECURE_ITEM_NONCE, ret);
			goto out;
		}

		if (isAllZero(vk_id, VK_ID_LEN)) {
			LOGE("%s: Failed to generate VK ID\n", __func__);
			ret = VK_ERR_TZ_API_CRYPTO;
			goto out;
		} else {
			for (i = 0; i != VK_ID_LEN; ++i) {
				switch (vk_id[i] % 3) {
					case 0:
						vk_id[i] = vk_id[i] % 10 + 48; // 0-9
						break;
					case 1:
						vk_id[i] = vk_id[i] % 26 + 65; // A-Z
						break;
					case 2:
						vk_id[i] = vk_id[i] % 26 + 97; // a-z
						break;
				}
			}

			LOGI("%s: Set the VaultKeeper ID : %s\n", __func__, vk_id);
			memcpy(vault->reserved + VEK_LEN + AP_SN_LEN, vk_id, VK_ID_LEN);
		}
	}
#ifdef __DEBUG__
	print_hex_to_string(vault->reserved + VEK_LEN + AP_SN_LEN, VK_ID_LEN);
#endif
	// <----6. Check & Generate VK ID

	ret = target->vault.write_unsheltered_vault(target, req, rsp, vault);
	if (ret != VK_SUCCESS) {
		LOGE("%s: Failed write vault(%d)\n", __func__, ret);
		goto out;
	}

	// Delivery ID to NWd to set ro-property
	memcpy(rsp->msg, vault->reserved + VEK_LEN + AP_SN_LEN, VK_ID_LEN);
	rsp->msg_len = VK_ID_LEN;

out:
	if (vault != NULL) {
		memset((uint8_t*)vault, 0, sizeof(unsvault_t));
		target->platform.memory_free(vault);
	}

	memset(vault_eKey, 0, VEK_LEN);
	memset(vk_id, 0, VK_ID_LEN);

	return ret;
}