// SPDX-License-Identifier: BSD-2-Clause
/*
 * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 */

#include <string.h>
#include <stdlib.h>
#include <types_ext.h>
#include <tee_api_types.h>
#include <trace.h>
#include <initcall.h>
#include <utee_defines.h>
#include <tegra_se_ccc.h>
#include <tegra_se_ccc_aes.h>
#include <tegra_se_ccc_errno.h>
#include <tegra_se_driver_aes.h>
#include <drivers/tegra/tegra_se_keyslot.h>

TEE_Result se_aes_encrypt(uint8_t *src, size_t src_len,
			  uint8_t *dst, size_t dst_len,
			  se_aes_op_mode mode,
			  uint32_t keyslot, uint8_t *iv)
{
	status_t ret = NO_ERROR;

	if (src == NULL || src_len == 0 || dst == NULL || dst_len == 0)
		return TEE_ERROR_BAD_PARAMETERS;

	ret = tegra_se_ccc_aes_encrypt(src, src_len,
				       dst, dst_len,
				       mode, keyslot,
				       TEGRA_SE_KEY_256_SIZE, iv);
	if (ret != NO_ERROR) {
		EMSG("%s failed with: %d", __func__, ret);
		return TEE_ERROR_GENERIC;
	}

        return TEE_SUCCESS;
}

TEE_Result se_aes_cmac(se_aes_op_mode mode,
		       se_aes_keyslot_t keyslot, uint32_t key_len,
		       uint8_t *src, uint32_t src_len,
		       uint8_t *mac,  uint32_t mac_len)
{
	status_t ret = NO_ERROR;

	if (src == NULL || src_len == 0)
		return TEE_ERROR_BAD_PARAMETERS;

	if (mac == NULL || mac_len != TEGRA_SE_AES_BLOCK_SIZE)
		return TEE_ERROR_BAD_PARAMETERS;


	ret = tegra_se_ccc_aes_encrypt(src, src_len,
				       mac, mac_len,
				       mode, keyslot,
				       key_len, NULL);
	if (ret != NO_ERROR) {
		EMSG("%s failed with: %d", __func__, ret);
		return TEE_ERROR_GENERIC;
	}

        return TEE_SUCCESS;
}

TEE_Result se_nist_sp_800_108_hmac_kdf_to_memory(se_aes_keyslot_t keyslot, uint32_t key_len,
						 char const *label, uint32_t label_len,
						 char const *context, uint32_t context_len,
						 uint32_t dk_len, uint8_t *out_dk)
{
	TEE_Result ret = TEE_SUCCESS;
	uint8_t zero_byte[] = { 0U };
	uint8_t *p, tmp_buf[TEE_SHA256_HASH_SIZE] = { 0U };
	uint32_t counter[] = { TEE_U32_TO_BIG_ENDIAN( 1 ) };
	uint32_t L[] = { TEE_U32_TO_BIG_ENDIAN(dk_len * 8) };
	uint32_t tmp_len;
	uint8_t *message = NULL;
	uint8_t *mptr;
	uint32_t msg_len;

	if ((key_len != TEGRA_SE_KEY_128_SIZE) &&
			(key_len != TEGRA_SE_KEY_256_SIZE))
		return TEE_ERROR_BAD_PARAMETERS;

	if ((dk_len % TEE_AES_BLOCK_SIZE) != 0)
		return TEE_ERROR_BAD_PARAMETERS;

	if (!label || !context || !out_dk)
		return TEE_ERROR_BAD_PARAMETERS;

	/*
	 *  Regarding to NIST-SP 800-108
	 *  message = counter || label || 0 || context || L
	 *
	 *  A || B = The concatenation of binary strings A and B.
	 */
	msg_len = sizeof(counter) + context_len + label_len + sizeof(L) +
		sizeof(zero_byte);
	message = calloc(1, msg_len);
	if (message == NULL)
		return TEE_ERROR_OUT_OF_MEMORY;

	/* Concatenate the messages */
	mptr = message;
	memcpy(mptr, counter, sizeof(counter));
	mptr += sizeof(counter);
	memcpy(mptr, label, label_len);
	mptr += label_len;
	memcpy(mptr, zero_byte, sizeof(zero_byte));
	mptr += sizeof(zero_byte);
	memcpy(mptr, context, context_len);
	mptr += context_len;
	memcpy(mptr, L, sizeof(L));

	tmp_len = dk_len;
	p = out_dk;
	while (tmp_len > 0) {
		if (tmp_len >= TEE_SHA256_HASH_SIZE) {
			ret = tegra_se_ccc_hmac_sha256(keyslot, key_len,
						       message, msg_len,
						       p);
			if (ret != NO_ERROR) {
				EMSG("HMAC-SHA256 failed with: %d", ret);
				ret = TEE_ERROR_GENERIC;
				goto kdf_error;
			}
			tmp_len -= TEE_SHA256_HASH_SIZE;
			p += TEE_SHA256_HASH_SIZE;
		} else {
			ret = tegra_se_ccc_hmac_sha256(keyslot, key_len,
						       message, msg_len,
						       tmp_buf);
			if (ret != NO_ERROR) {
				EMSG("HMAC-SHA256 failed with: %d", ret);
				ret = TEE_ERROR_GENERIC;
				goto kdf_error;
			}
			memcpy(p, tmp_buf, tmp_len);
			break;
		}
		message[0] += 1;
	}

kdf_error:
	free(message);
	return ret;
}

TEE_Result se_nist_sp_800_108_hmac_kdf_to_keyslot(uint32_t source_keyslot,
						  const char *label, uint32_t label_len,
						  const char *context, uint32_t context_len,
						  uint32_t manifest, uint32_t target_keyslot)
{
	return tegra_se_ccc_kdf_encoder(source_keyslot,
					label, label_len,
					context, context_len,
					manifest, target_keyslot);

}

TEE_Result se_clear_aes_keyslots(void)
{
	status_t ret = NO_ERROR;
	uint32_t i;

	for (i = SE_AES_KEYSLOT_10; i <= SE_AES_KEYSLOT_13; i++) {
		ret = tegra_se_ccc_clear_aes_keyslot(i);
		if (ret != NO_ERROR) {
			EMSG("%s failed with: %d", __func__, ret);
			return TEE_ERROR_GENERIC;
		}
	}

	return TEE_SUCCESS;
}
