/*
 * Copyright (c) 2025 Nordic Semiconductor ASA
 *
 * SPDX-License-Identifier: Apache-2.0
 */

#include "mcuboot_config/mcuboot_config.h"

#if defined(MCUBOOT_USE_PSA_CRYPTO)

#include <stddef.h>
#include <inttypes.h>
#include <string.h>

/* We are not really using the MBEDTLS but need the ASN.1 parsing functions */
#define MBEDTLS_ASN1_PARSE_C

#include "bootutil/crypto/sha.h"
#include "mbedtls/build_info.h"
#include "mbedtls/oid.h"
#include "mbedtls/asn1.h"

/*
 * In Mbed TLS 4.x, MBEDTLS_OID_EC_ALG_UNRESTRICTED and
 * MBEDTLS_OID_EC_GRP_SECP256R1 were moved from the public
 * `mbedtls/oid.h` to the private `tf-psa-crypto/utilities/crypto_oid.h`.
 * Rather than pull a private header, define the raw OID byte strings
 * locally when the public macros are not available. Values are taken
 * from X9.62 / SEC2:
 *   id-ecPublicKey (1.2.840.10045.2.1)       → {0x2a,0x86,0x48,0xce,0x3d,0x02,0x01}
 *   secp256r1      (1.2.840.10045.3.1.7)     → {0x2a,0x86,0x48,0xce,0x3d,0x03,0x01,0x07}
 */
#if !defined(MBEDTLS_OID_EC_ALG_UNRESTRICTED)
#define MBEDTLS_OID_EC_ALG_UNRESTRICTED "\x2a\x86\x48\xce\x3d\x02\x01"
#endif
#if !defined(MBEDTLS_OID_EC_GRP_SECP256R1)
#define MBEDTLS_OID_EC_GRP_SECP256R1    "\x2a\x86\x48\xce\x3d\x03\x01\x07"
#endif

#include "bootutil/image.h"
#include "bootutil/enc_key.h"
#include "bootutil/sign_key.h"
#include "bootutil/crypto/common.h"

#include "bootutil_priv.h"
#include "bootutil/bootutil_log.h"

BOOT_LOG_MODULE_REGISTER(mcuboot_psa_enc);

#if defined(MCUBOOT_HMAC_SHA512)
#define PSA_HMAC_HKDF_SHA PSA_ALG_SHA_512
#else
#define PSA_HMAC_HKDF_SHA PSA_ALG_SHA_256
#endif

#if defined(MCUBOOT_ENCRYPT_EC256)
#define NUM_ECC_BYTES (256 / 8)
static const uint8_t ec_pubkey_oid[] = MBEDTLS_OID_EC_ALG_UNRESTRICTED;
static const uint8_t ec_secp256r1_oid[] = MBEDTLS_OID_EC_GRP_SECP256R1;
#define ECC_FAMILY PSA_ECC_FAMILY_SECP_R1
#endif /* defined(MCUBOOT_ENCRYPT_EC256) */
#if defined(MCUBOOT_ENCRYPT_X25519)
#define X25519_OID "\x6e"
static const uint8_t ec_pubkey_oid[] = MBEDTLS_OID_ISO_IDENTIFIED_ORG \
                                       MBEDTLS_OID_ORG_GOV X25519_OID;
#define ECC_FAMILY PSA_ECC_FAMILY_MONTGOMERY
#endif /* defined(MCUBOOT_ENCRYPT_X25519) */

/* Partitioning of HKDF derived material, from the exchange derived key */
/* AES key encryption key */
#define HKDF_AES_KEY_INDEX  0
#define HKDF_AES_KEY_SIZE   (BOOT_ENC_KEY_SIZE)
/* MAC feed */
#define HKDF_MAC_FEED_INDEX (HKDF_AES_KEY_INDEX + HKDF_AES_KEY_SIZE)
#if !defined(MCUBOOT_HMAC_SHA512)
#define HKDF_MAC_FEED_SIZE  (32)
#else
#define HKDF_MAC_FEED_SIZE  (64)
#endif
/* Total size */
#define HKDF_SIZE           (HKDF_AES_KEY_SIZE + HKDF_MAC_FEED_SIZE)

#if defined(MCUBOOT_ENCRYPT_EC256)
/* Fixme: This duplicates code from encrypted.c and depends on mbedtls */

/*
 * Parses the output of `imgtool keygen`, which produces a PKCS#8 elliptic
 * curve keypair. See RFC5208 and RFC5915.
 */
static int
parse_priv_enckey(uint8_t **p, uint8_t *end, uint8_t *private_key)
{
    size_t len;
    int version;
    mbedtls_asn1_buf alg;
    mbedtls_asn1_buf param;

    if (mbedtls_asn1_get_tag(p, end, &len,
                             MBEDTLS_ASN1_CONSTRUCTED | MBEDTLS_ASN1_SEQUENCE) != 0) {
        return -1;
    }

    if (*p + len != end) {
        return -1;
    }

    version = 0;
    if (mbedtls_asn1_get_int(p, end, &version) || version != 0) {
        return -1;
    }

    if (mbedtls_asn1_get_alg(p, end, &alg, &param) != 0) {
        return -1;
    }

    if (alg.ASN1_CONTEXT_MEMBER(len) != sizeof(ec_pubkey_oid) - 1 ||
        memcmp(alg.ASN1_CONTEXT_MEMBER(p), ec_pubkey_oid, sizeof(ec_pubkey_oid) - 1)) {
        return -1;
    }
    if (param.ASN1_CONTEXT_MEMBER(len) != sizeof(ec_secp256r1_oid) - 1 ||
        memcmp(param.ASN1_CONTEXT_MEMBER(p), ec_secp256r1_oid, sizeof(ec_secp256r1_oid) - 1)) {
        return -1;
    }

    if (mbedtls_asn1_get_tag(p, end, &len, MBEDTLS_ASN1_OCTET_STRING) != 0) {
        return -1;
    }

    /* RFC5915 - ECPrivateKey */

    if (mbedtls_asn1_get_tag(p, end, &len,
                             MBEDTLS_ASN1_CONSTRUCTED | MBEDTLS_ASN1_SEQUENCE) != 0) {
        return -1;
    }

    version = 0;
    if (mbedtls_asn1_get_int(p, end, &version) || version != 1) {
        return -1;
    }

    /* privateKey */

    if (mbedtls_asn1_get_tag(p, end, &len, MBEDTLS_ASN1_OCTET_STRING) != 0) {
        return -1;
    }

    if (len != NUM_ECC_BYTES) {
        return -1;
    }

    memcpy(private_key, *p, len);

    /* publicKey usually follows but is not parsed here */

    return 0;
}
#endif /* defined(MCUBOOT_ENCRYPT_EC256) */

#if defined(MCUBOOT_ENCRYPT_X25519)
/* Fixme: This duplicates code from encrypted.c and depends on mbedtls */
static int
parse_priv_enckey(uint8_t **p, uint8_t *end, uint8_t *private_key)
{
    size_t len;
    int version;
    mbedtls_asn1_buf alg;
    mbedtls_asn1_buf param;

    if (mbedtls_asn1_get_tag(p, end, &len, MBEDTLS_ASN1_CONSTRUCTED |
                                           MBEDTLS_ASN1_SEQUENCE) != 0) {
        return -1;
    }

    if (*p + len != end) {
        return -2;
    }

    version = 0;
    if (mbedtls_asn1_get_int(p, end, &version) || version != 0) {
        return -3;
    }

    if (mbedtls_asn1_get_alg(p, end, &alg, &param) != 0) {
        return -4;
    }

    if (alg.ASN1_CONTEXT_MEMBER(len) != sizeof(ec_pubkey_oid) - 1 ||
        memcmp(alg.ASN1_CONTEXT_MEMBER(p), ec_pubkey_oid, sizeof(ec_pubkey_oid) - 1)) {
        return -5;
    }

    if (mbedtls_asn1_get_tag(p, end, &len, MBEDTLS_ASN1_OCTET_STRING) != 0) {
        return -6;
    }

    if (mbedtls_asn1_get_tag(p, end, &len, MBEDTLS_ASN1_OCTET_STRING) != 0) {
        return -7;
    }

    if (len != EC_PRIVK_LEN) {
        return -8;
    }

    memcpy(private_key, *p, EC_PRIVK_LEN);
    return 0;
}
#endif /* defined(MCUBOOT_ENCRYPT_X25519) */

void bootutil_aes_ctr_init(bootutil_aes_ctr_context *ctx)
{
    psa_status_t psa_ret = psa_crypto_init();

    if (psa_ret != PSA_SUCCESS) {
        BOOT_LOG_ERR("AES init PSA crypto init failed %d", psa_ret);
        assert(0);
    }

    ctx->key = PSA_KEY_ID_NULL;
}

void bootutil_aes_ctr_drop(bootutil_aes_ctr_context *ctx)
{
    psa_status_t psa_ret = psa_destroy_key(ctx->key);

    if (psa_ret != PSA_SUCCESS) {
        BOOT_LOG_WRN("aes_ctr_drop: destruction failed %d", psa_ret);
        /* This should never happen. If we fail to destroy key this happens
         * either because it is invalid key number or something is really
         * wrong; either way we have no way to recover.
         */
        assert(0);
    }

    ctx->key = PSA_KEY_ID_NULL;
}

int bootutil_aes_ctr_set_key(bootutil_aes_ctr_context *ctx, const uint8_t *k)
{
    psa_status_t psa_ret = PSA_ERROR_BAD_STATE;
    psa_key_attributes_t kattr = PSA_KEY_ATTRIBUTES_INIT;

    psa_set_key_type(&kattr, PSA_KEY_TYPE_AES);
    psa_set_key_usage_flags(&kattr, PSA_KEY_USAGE_DECRYPT | PSA_KEY_USAGE_ENCRYPT);
    psa_set_key_algorithm(&kattr, PSA_ALG_CTR);

    psa_ret = psa_import_key(&kattr, k, HKDF_AES_KEY_SIZE, &ctx->key);
    if (psa_ret != PSA_SUCCESS) {
        BOOT_LOG_ERR("aes_ctr_set_key; import failed %d", psa_ret);
        return -1;
    }
    return 0;
}

#if defined(MCUBOOT_ENC_IMAGES)
extern const struct bootutil_key bootutil_enc_key;
/*
 * Decrypt an encryption key TLV.
 *
 * @param buf An encryption TLV read from flash (build time fixed length)
 * @param enckey An AES-128 or AES-256 key sized buffer to store to plain key.
 */
int
boot_decrypt_key(const uint8_t *buf, uint8_t *enckey)
{
    uint8_t derived_key[HKDF_SIZE];
    uint8_t *cp;
    uint8_t *cpend;
    uint8_t private_key[EC_PRIVK_LEN];
    size_t len;
    psa_status_t psa_ret = PSA_ERROR_BAD_STATE;
    psa_status_t psa_cleanup_ret = PSA_ERROR_BAD_STATE;
    psa_key_id_t kid;
    psa_key_attributes_t kattr = PSA_KEY_ATTRIBUTES_INIT;
    psa_key_derivation_operation_t key_do = PSA_KEY_DERIVATION_OPERATION_INIT;
    psa_algorithm_t key_do_alg;
    int rc = -1;

    cp = (uint8_t *)bootutil_enc_key.key;
    cpend = cp + *bootutil_enc_key.len;

    /* The psa_cipher_decrypt needs initialization vector of proper length at
     * the beginning of the input buffer.
     */
    uint8_t iv_and_key[PSA_CIPHER_IV_LENGTH(PSA_KEY_TYPE_AES, PSA_ALG_CTR) +
                       BOOT_ENC_KEY_SIZE];

    BOOT_LOG_DBG("boot_decrypt_key: PSA ED25519");

    psa_ret = psa_crypto_init();
    if (psa_ret != PSA_SUCCESS) {
        BOOT_LOG_ERR("PSA crypto init failed %d", psa_ret);
        return -1;
    }

    /*
     * * Load the stored decryption private key
     */
    rc = parse_priv_enckey(&cp, cpend, private_key);
    if (rc) {
        BOOT_LOG_ERR("Failed to parse ASN1 private key");
        return rc;
    }

    psa_set_key_type(&kattr, PSA_KEY_TYPE_ECC_KEY_PAIR(ECC_FAMILY));
    psa_set_key_usage_flags(&kattr, PSA_KEY_USAGE_DERIVE);
    psa_set_key_algorithm(&kattr, PSA_ALG_ECDH);

    psa_ret = psa_import_key(&kattr, private_key, sizeof(private_key), &kid);
    memset(private_key, 0, sizeof(private_key));
    psa_reset_key_attributes(&kattr);
    if (psa_ret != PSA_SUCCESS) {
        BOOT_LOG_ERR("Built-in key import failed %d", psa_ret);
        return -1;
    }

    key_do_alg = PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_HMAC_HKDF_SHA));

    psa_ret = psa_key_derivation_setup(&key_do, key_do_alg);
    if (psa_ret != PSA_SUCCESS) {
        psa_cleanup_ret = psa_destroy_key(kid);
        if (psa_cleanup_ret != PSA_SUCCESS) {
            BOOT_LOG_WRN("Built-in key destruction failed %d", psa_cleanup_ret);
        }
        BOOT_LOG_ERR("Key derivation setup failed %d", psa_ret);
        return -1;
    }

    /* Note: PSA 1.1.2 does not have psa_key_agreement that would be useful here
     * as it could just add the derived key to the storage and return key id.
     * Instead, we have to use the code below to generate derived key and put it
     * into storage, to obtain the key id we can then use with psa_mac_* functions.
     */
    psa_ret = psa_key_derivation_key_agreement(&key_do, PSA_KEY_DERIVATION_INPUT_SECRET,
                                               kid, &buf[EC_PUBK_INDEX],
                                               EC_PUBK_LEN);
    psa_cleanup_ret = psa_destroy_key(kid);
    if (psa_cleanup_ret != PSA_SUCCESS) {
        BOOT_LOG_WRN("Built-in key destruction failed %d", psa_cleanup_ret);
    }
    if (psa_ret != PSA_SUCCESS) {
        psa_cleanup_ret = psa_key_derivation_abort(&key_do);
        if (psa_cleanup_ret != PSA_SUCCESS) {
            BOOT_LOG_WRN("Key derivation abort failed %d", psa_ret);
        }

        BOOT_LOG_ERR("Key derivation failed %d", psa_ret);
        return -1;
    }

    /* Only info, no salt */
    psa_ret = psa_key_derivation_input_bytes(&key_do, PSA_KEY_DERIVATION_INPUT_INFO,
                                             (const uint8_t *)"MCUBoot_ECIES_v1", 16);
    if (psa_ret != PSA_SUCCESS) {
        psa_cleanup_ret = psa_key_derivation_abort(&key_do);
        if (psa_cleanup_ret != PSA_SUCCESS) {
            BOOT_LOG_WRN("Key derivation abort failed %d", psa_ret);
        }
        BOOT_LOG_ERR("Key derivation failed %d", psa_ret);
        return -1;
    }

    len = HKDF_SIZE;
    psa_ret = psa_key_derivation_output_bytes(&key_do, derived_key, len);
    psa_cleanup_ret = psa_key_derivation_abort(&key_do);
    if (psa_cleanup_ret != PSA_SUCCESS) {
        BOOT_LOG_WRN("Key derivation cleanup failed %d", psa_ret);
    }
    if (psa_ret != PSA_SUCCESS) {
        BOOT_LOG_ERR("Key derivation failed %d", psa_ret);
        return -1;
    }

    /* The derived key consists of BOOT_ENC_KEY_SIZE bytes
     * followed by BOOTUTIL_CRYPTO_SHA256_DIGEST_SIZE bytes. Both parts will
     * be imported at the point where needed and discarded immediately after.
     */
    psa_set_key_type(&kattr, PSA_KEY_TYPE_HMAC);
    psa_set_key_usage_flags(&kattr, PSA_KEY_USAGE_VERIFY_MESSAGE);
    psa_set_key_algorithm(&kattr, PSA_ALG_HMAC(PSA_HMAC_HKDF_SHA));

    /* Import the MAC tag key part of derived key */
    psa_ret = psa_import_key(&kattr,
                             &derived_key[HKDF_MAC_FEED_INDEX],
                             HKDF_MAC_FEED_SIZE, &kid);
    psa_reset_key_attributes(&kattr);
    if (psa_ret != PSA_SUCCESS) {
        memset(derived_key, 0, sizeof(derived_key));
        BOOT_LOG_ERR("MAC key import failed %d", psa_ret);
        return -1;
    }

    /* Verify the MAC tag of the random encryption key */
    psa_ret = psa_mac_verify(kid, PSA_ALG_HMAC(PSA_HMAC_HKDF_SHA),
                             &buf[EC_CIPHERKEY_INDEX], EC_CIPHERKEY_LEN,
                             &buf[EC_TAG_INDEX],
                             EC_TAG_LEN);
    psa_cleanup_ret = psa_destroy_key(kid);
    if (psa_cleanup_ret != PSA_SUCCESS) {
        BOOT_LOG_WRN("MAC key destruction failed %d", psa_cleanup_ret);
    }
    if (psa_ret != PSA_SUCCESS) {
        memset(derived_key, 0, sizeof(derived_key));
        BOOT_LOG_ERR("MAC verification failed %d", psa_ret);
        return -1;
    }

    /* The derived key is used in AES decryption of random key */
    psa_set_key_type(&kattr, PSA_KEY_TYPE_AES);
    psa_set_key_usage_flags(&kattr, PSA_KEY_USAGE_DECRYPT);
    psa_set_key_algorithm(&kattr, PSA_ALG_CTR);

    /* Import the AES partition of derived key, the first 16 bytes */
    psa_ret = psa_import_key(&kattr, &derived_key[HKDF_AES_KEY_INDEX],
                             HKDF_AES_KEY_SIZE, &kid);
    memset(derived_key, 0, sizeof(derived_key));
    if (psa_ret != PSA_SUCCESS) {
        BOOT_LOG_ERR("AES key import failed %d", psa_ret);
        return -1;
    }

    /* Decrypt the random AES encryption key with AES and the key obtained
     * at derivation. */
    memset(&iv_and_key[0], 0, PSA_CIPHER_IV_LENGTH(PSA_KEY_TYPE_AES, PSA_ALG_CTR));
    memcpy(&iv_and_key[PSA_CIPHER_IV_LENGTH(PSA_KEY_TYPE_AES, PSA_ALG_CTR)],
           &buf[EC_CIPHERKEY_INDEX],
	   sizeof(iv_and_key) - PSA_CIPHER_IV_LENGTH(PSA_KEY_TYPE_AES, PSA_ALG_CTR));

    len = 0;
    psa_ret = psa_cipher_decrypt(kid, PSA_ALG_CTR, iv_and_key, sizeof(iv_and_key),
                                 enckey, BOOT_ENC_KEY_SIZE, &len);
    memset(iv_and_key, 0, sizeof(iv_and_key));
    psa_cleanup_ret = psa_destroy_key(kid);
    if (psa_cleanup_ret != PSA_SUCCESS) {
	BOOT_LOG_WRN("AES key destruction failed %d", psa_cleanup_ret);
    }
    if (psa_ret != PSA_SUCCESS || len != BOOT_ENC_KEY_SIZE) {
        memset(enckey, 0, BOOT_ENC_KEY_SIZE);
        BOOT_LOG_ERR("Random key decryption failed %d", psa_ret);
        return -1;
    }

    return 0;
}

int bootutil_aes_ctr_encrypt(bootutil_aes_ctr_context *ctx, uint8_t *counter,
        const uint8_t *m, uint32_t mlen, size_t blk_off, uint8_t *c)
{
    int ret = 0;
    psa_status_t psa_ret = PSA_ERROR_BAD_STATE;
    const psa_key_id_t kid = ctx->key;
    psa_cipher_operation_t psa_op;
    size_t elen = 0;	/* Decrypted length */

    /* PSA cipher API handles CTR block alignment internally. */
    (void)blk_off;

    /* Fixme: calling psa_crypto_init multiple times is not a problem,
     * yet the code here is only present because there is not general
     * crypto init. */
    psa_ret = psa_crypto_init();
    if (psa_ret != PSA_SUCCESS) {
        BOOT_LOG_ERR("PSA crypto init failed %d", psa_ret);
        ret = -1;
        goto gone;
    }

    psa_op = psa_cipher_operation_init();

    /* This could be done with psa_cipher_decrypt one-shot operation, but
     * multi-part operation is used to avoid re-allocating input buffer
     * to account for IV in front of data.
     */
    psa_ret = psa_cipher_encrypt_setup(&psa_op, kid, PSA_ALG_CTR);
    if (psa_ret != PSA_SUCCESS) {
        BOOT_LOG_ERR("AES enc setup failed %d", psa_ret);
        ret = -1;
        goto gone;
    }

    /* Fixme: hardcoded counter  size, but it is hardcoded everywhere */
    psa_ret = psa_cipher_set_iv(&psa_op, counter, 16);
    if (psa_ret != PSA_SUCCESS) {
	BOOT_LOG_ERR("AES enc IV set failed %d", psa_ret);
        ret = -1;
        goto gone_after_setup;
    }

    psa_ret = psa_cipher_update(&psa_op, m, mlen, c, mlen, &elen);
    if (psa_ret != PSA_SUCCESS) {
	BOOT_LOG_ERR("AES enc encryption failed %d", psa_ret);
        ret = -1;
        goto gone_after_setup;
    }

gone_after_setup:
    psa_ret = psa_cipher_abort(&psa_op);
    if (psa_ret != PSA_SUCCESS) {
        BOOT_LOG_WRN("AES enc cipher abort failed %d", psa_ret);
        /* Intentionally not changing the ret */
    }
gone:
    return ret;
}

int bootutil_aes_ctr_decrypt(bootutil_aes_ctr_context *ctx, uint8_t *counter,
        const uint8_t *c, uint32_t clen, size_t blk_off, uint8_t *m)
{
    int ret = 0;
    psa_status_t psa_ret = PSA_ERROR_BAD_STATE;
    const psa_key_id_t kid = ctx->key;
    psa_cipher_operation_t psa_op;
    size_t dlen = 0;	/* Decrypted length */

    /* PSA cipher API handles CTR block alignment internally. */
    (void)blk_off;

    /* Fixme: the init should already happen before calling the function, but
     * somehow it does not, for example when recovering in swap.
     */
    psa_ret = psa_crypto_init();
    if (psa_ret != PSA_SUCCESS) {
        BOOT_LOG_ERR("PSA crypto init failed %d", psa_ret);
        ret = -1;
        goto gone;
    }

    psa_op = psa_cipher_operation_init();

    /* This could be done with psa_cipher_decrypt one-shot operation, but
     * multi-part operation is used to avoid re-allocating input buffer
     * to account for IV in front of data.
     */
    psa_ret = psa_cipher_decrypt_setup(&psa_op, kid, PSA_ALG_CTR);
    if (psa_ret != PSA_SUCCESS) {
        BOOT_LOG_ERR("AES dec setup failed %d", psa_ret);
        ret = -1;
        goto gone;
    }

    /* Fixme: hardcoded counter  size, but it is hardcoded everywhere */
    psa_ret = psa_cipher_set_iv(&psa_op, counter, 16);
    if (psa_ret != PSA_SUCCESS) {
	BOOT_LOG_ERR("AES dec IV set failed %d", psa_ret);
        ret = -1;
        goto gone_after_setup;
    }

    psa_ret = psa_cipher_update(&psa_op, c, clen, m, clen, &dlen);
    if (psa_ret != PSA_SUCCESS) {
	BOOT_LOG_ERR("AES dec decryption failed %d", psa_ret);
        ret = -1;
        goto gone_after_setup;
    }

gone_after_setup:
    psa_ret = psa_cipher_abort(&psa_op);
    if (psa_ret != PSA_SUCCESS) {
        BOOT_LOG_WRN("PSA dec abort failed %d", psa_ret);
        /* Intentionally not changing the ret */
    }
gone:
    return ret;
}
#endif /* defined(MCUBOOT_ENC_IMAGES) */

#endif /* defined(MCUBOOT_USE_PSA_CRYPTO) */