#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_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_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_RPMB_OPEN_PART;
        }

    } else if (ret != QSEE_STOR_SUCCESS) {
        LOGE("%s: Failed to open partition(%d)\n", __func__, ret);
        return EM_ERR_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_RPMB_READ_SECTOR;
		goto out;
	}
	goto out;
    }

    if (em_is_all_zero(data, EM_RPMB_SINGLE_BLOCK_SIZE*2) == EM_SUCCESS) {
        LOGE("%s: The block is all zero\n", __func__);
        ret = EM_ERR_ALL_ZERO;
        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_RPMB_WRITE_SECTOR;
		goto out;
	}
	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_SIZE_KEY] = {};
    uint8_t iv[EM_SIZE_IV] = {};
    uint32_t plaintext_len = 0;
    uint32_t encrypt_len = 0;
    uint32_t unencrypt_len = 0;

    // Param check
    if (buf == NULL) {
        LOGE("%s : buf is NULL\n", __func__);
        ret = EM_ERR_INVALID_ARG;
        goto out;
    }

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

    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
    ret = em_qsee_rpmb_read(0, core);
    if (ret == EM_ERR_ALL_ZERO) {
        LOGI("%s : core data is all zero\n", __func__);
        goto out;
    }

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

    // Get Key
    ret = em_crypto_kdf(key, EM_SIZE_KEY, iv, EM_SIZE_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_SIZE_MAGIC + EM_SIZE_GCM_TAG;
    encrypt_len = EM_RPMB_SINGLE_BLOCK_SIZE - unencrypt_len;

    // Unwap data
    ret = em_crypto_aes_256_gcm_decrypt(core + unencrypt_len, encrypt_len,
                                        buf + unencrypt_len, &plaintext_len,
                                        key, EM_SIZE_KEY,
                                        iv, EM_SIZE_IV,
                                        core + EM_SIZE_MAGIC, EM_SIZE_GCM_TAG);
    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_INTERNAL;
        goto out;
    }

out:
    memset(key, 0, EM_SIZE_KEY);
    memset(iv, 0, EM_SIZE_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_SIZE_KEY] = {};
    uint8_t iv[EM_SIZE_IV] = {};
    uint32_t ciphertext_len = 0;
    uint32_t encrypt_len = 0;
    uint32_t unencrypt_len = 0;

    // Param check
    if (buf == NULL) {
        LOGE("%s : buf is NULL\n", __func__);
        ret = EM_ERR_INVALID_ARG;
        goto out;
    }

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

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

    // Get Key
    ret = em_crypto_kdf(key, EM_SIZE_KEY, iv, EM_SIZE_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_SIZE_MAGIC + EM_SIZE_GCM_TAG;
    encrypt_len = EM_RPMB_SINGLE_BLOCK_SIZE - unencrypt_len;

    // Wrap data
    ret = em_crypto_aes_256_gcm_encrypt(buf + unencrypt_len, encrypt_len,
                                        core + unencrypt_len, &ciphertext_len,
                                        key, EM_SIZE_KEY,
                                        iv, EM_SIZE_IV,
                                        core + EM_SIZE_MAGIC, EM_SIZE_GCM_TAG);
    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_INTERNAL;
        goto out;
    }

    // Read from RPMB
    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_SIZE_KEY);
    memset(iv, 0, EM_SIZE_IV);

    return ret;
}

int em_check_provision()
{
	int ret;
	qsee_secctrl_secure_status_t status;

	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_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_RPMB_NOT_PROVISION;
		goto out;
	}

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