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

#include <stdio.h>
#include <string.h>
#include <trace.h>
#include <config.h>
#include <crypto/crypto.h>
#include <drivers/tegra/tegra_fuse.h>
#include <drivers/tegra/tegra_se_aes.h>
#include <drivers/tegra/tegra_se_kdf.h>
#include "drivers/tegra/tegra_se_keyslot.h"
#include "jetson_ekb_extraction_t264.h"

#define EKB_RK_MANIFEST KEYSIZE_256_TZ_KDK_MANIFEST
#define EKB_AK_MANIFEST KEYSIZE_256_TZ_CMAC_MANIFEST
#define EKB_EK_MANIFEST KEYSIZE_256_TZ_ENC_MANIFEST

TEE_Result jetson_derive_keys(void)
{
	TEE_Result rc = TEE_SUCCESS;

	rc = tegra_se_nist_sp_800_108_hmac_kdf_to_ks(SE_AES_KEYSLOT_TZ_RK,
						    (char *)"ekb", strlen("ekb"),
						    (char *)"root", strlen("root"),
						    EKB_RK_MANIFEST,
						    SE_AES_KEYSLOT_EKB_RK);
	if (rc != TEE_SUCCESS) {
		EMSG("%s: Failed to derive ekb root key (%x)", __func__, rc);
		goto key_derivation_fail;
	}

	rc = tegra_se_nist_sp_800_108_hmac_kdf_to_ks(SE_AES_KEYSLOT_EKB_RK,
						    (char *)"ekb", strlen("ekb"),
						    (char *)"encryption", strlen("encryption"),
						    EKB_EK_MANIFEST,
						    SE_AES_KEYSLOT_EKB_EK);
	if (rc != TEE_SUCCESS) {
		EMSG("%s: Failed to derive ekb encryption key (%x)", __func__, rc);
		goto key_derivation_fail;
	}

	rc = tegra_se_nist_sp_800_108_hmac_kdf_to_ks(SE_AES_KEYSLOT_EKB_RK,
						    (char *)"ekb", strlen("ekb"),
						    (char *)"authentication", strlen("authentication"),
						    EKB_AK_MANIFEST,
						    SE_AES_KEYSLOT_EKB_AK);
	if (rc != TEE_SUCCESS) {
		EMSG("%s: Failed to derive ekb authentication key (%x)", __func__, rc);
		goto key_derivation_fail;
	}

key_derivation_fail:
	return rc;
}

TEE_Result jetson_ekb_authenticate_mac(struct ekb_content *ekb)
{
	TEE_Result rc = TEE_SUCCESS;
	struct ekb_content_header *content_header;
	uint8_t ekb_cmac_verify[TEGRA_SE_AES_BLOCK_SIZE] = {0};

	if (!ekb)
		return TEE_ERROR_BAD_PARAMETERS;

	content_header = &ekb->content_header;

	rc = tegra_se_aes_cmac(SE_AES_KEYSLOT_EKB_AK, TEGRA_SE_KEY_256_SIZE,
			       (uint8_t *)content_header,
			       sizeof(struct ekb_content_header) + sizeof(ekb->random_iv)
			       + content_header->enc_data_len,
			       ekb_cmac_verify,  TEGRA_SE_AES_BLOCK_SIZE);
	if (rc) {
		EMSG("Calculate cmac failed with rc = 0x%08x", rc);
		return rc;
	}

	/* Verify the CMAC */
	if (memcmp(ekb->ekb_cmac, ekb_cmac_verify, TEGRA_SE_AES_BLOCK_SIZE)) {
		EMSG("Verify cmac failed");
		rc = TEE_ERROR_SECURITY;
	}

	return rc;
}

TEE_Result jetson_ekb_decrypt_content(struct ekb_content *ekb, uint8_t *decrypted_ekb_content)
{
	TEE_Result rc = TEE_SUCCESS;
	struct ekb_content_header *content_header;

	if (!ekb || !decrypted_ekb_content)
		return TEE_ERROR_BAD_PARAMETERS;

	content_header = &ekb->content_header;

	rc = tegra_se_aes_encrypt_cbc(ekb->ekb_ciphertext,
				      content_header->enc_data_len,
				      decrypted_ekb_content,
				      SE_AES_KEYSLOT_EKB_EK,
				      ekb->random_iv);
	if (rc)
		EMSG("Decrypte ekb content failed wit rc = 0x%08x", rc);

	return rc;
}
