/**
  ******************************************************************************
  * @file           : sm_crypto_mcu.c
  * @brief          : MCU cryptographic functions using hardware accelerators
  ******************************************************************************
  * @attention
  *
  * Copyright (c) 2025 STMicroelectronics.
  * All rights reserved.
  *
  * This software is licensed under terms that can be found in the LICENSE file
  * in the root directory of this software component.
  * If no LICENSE file comes with this software, it is provided AS-IS.
  *
  ******************************************************************************
  */

/* Includes ----------------------------------------------------------*/
#include "sm_crypto_mcu.h"
#include <string.h>
#include <stdlib.h>

#ifdef STM32WBA65xx
#include "hw.h"
#endif //STM32WBA65xx

#define USE_MCU_HW_RNG
#define USE_MCU_HW_ASYM

#define HASH_TIMEOUT ((uint32_t)5)
#define KDH_INFO_MAX_LEN_BYTES ((uint32_t)32)
#define SHA256_BYTE_SIZE ((uint32_t)32)





static HASH_HandleTypeDef *HHash = NULL;
static CRYP_HandleTypeDef *HCryp = NULL;

#ifdef USE_MCU_HW_RNG
static RNG_HandleTypeDef *HRng = NULL;
#endif // USE_MCU_HW_RNG





#if defined(STM32U545xx) || defined(STM32U585xx)
#include "stm32u5xx_hal.h"
#define HHASH_INIT_DATATYPE                     HASH_DATATYPE_8B
#define SM_HASH_SHA256_START                    HAL_HASHEx_SHA256_Start
#define SM_HASH_SHA256_ACCUMULATE(H, B, S, T)   HAL_HASHEx_SHA256_Accmlt(H, B, S)
#define SM_HASH_SHA256_ACCUMULATE_LAST          HAL_HASHEx_SHA256_Accmlt_End
#define SM_HASH_SHA1_START                      HAL_HASH_SHA1_Start
#define HASH_SET_ALGORITHM(I, A)

#define SM_HMAC_SHA256_START    HAL_HMACEx_SHA256_Start

#elif defined(STM32U385xx)
#include "stm32u3xx_hal.h"
#define HHASH_INIT_DATATYPE                     HASH_DATATYPE_8B
#define SM_HASH_SHA256_START                    HAL_HASHEx_SHA256_Start
#define SM_HASH_SHA256_ACCUMULATE(H, B, S, T)   HAL_HASHEx_SHA256_Accmlt(H, B, S)
#define SM_HASH_SHA256_ACCUMULATE_LAST          HAL_HASHEx_SHA256_Accmlt_End
#define SM_HASH_SHA1_START                      HAL_HASH_SHA1_Start
#define HASH_SET_ALGORITHM(I, A)

#define SM_HMAC_SHA256_START    HAL_HMACEx_SHA256_Start

#elif defined(STM32WBA52xx) || defined(STM32WBA55xx) || defined(STM32WBA65xx)
#include "stm32wbaxx_hal.h"
#define HHASH_INIT_DATATYPE     HASH_CR_DATATYPE_1

#define SM_HASH_SHA1_START(hhash, pInBuffer, Size, pOutBuffer, Timeout)         HAL_HASH_Start(hhash, pInBuffer, Size, pOutBuffer, Timeout)

#define SM_HASH_SHA256_START(hhash, pInBuffer, Size, pOutBuffer, Timeout)       HAL_HASH_Start(hhash, pInBuffer, Size, pOutBuffer, Timeout)
#define SM_HMAC_SHA256_START(hhash, pInBuffer, Size, pOutBuffer, Timeout)       HAL_HASH_HMAC_Start(hhash, pInBuffer, Size, pOutBuffer, Timeout)

#define SM_HASH_SHA256_ACCUMULATE(H, B, S, T)   HAL_HASH_Accumulate(H, B, S, T)
#define SM_HASH_SHA256_ACCUMULATE_LAST          HAL_HASH_AccumulateLast
#define HASH_SET_ALGORITHM(I, A)                I.Algorithm = A

#else
#error "product not supported"
#endif

#ifdef USE_MCU_HW_ASYM
static PKA_HandleTypeDef *HPka = NULL;
#define PKA_TIMEOUT ((uint32_t)3000)
#define PRIVATE_KEY_SIZE_BYTES ((uint32_t)32)
#define MODULUS_SIZE_BYTES PRIVATE_KEY_SIZE_BYTES
#define MODULUS_SIZE_WORDS ((MODULUS_SIZE_BYTES + 3)/4)
#define SECP256R1_ORDER_LEN_BITS ((uint32_t)256)
#define SECP256R1_A_SIGN ((uint32_t)1)
static const uint8_t SECP256R1_Prime[] =
{ 0xff, 0xff, 0xff, 0xff, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xff, 0xff, 0xff, 0xff, 0xff,
    0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff };
static const uint8_t SECP256R1_AbsA[] =
{ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
    0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03 };
static const uint8_t SECP256R1_B[] =
{ 0x5a, 0xc6, 0x35, 0xd8, 0xaa, 0x3a, 0x93, 0xe7, 0xb3, 0xeb, 0xbd, 0x55, 0x76,
    0x98, 0x86, 0xbc, 0x65, 0x1d, 0x06, 0xb0, 0xcc, 0x53, 0xb0, 0xf6, 0x3b,
    0xce, 0x3c, 0x3e, 0x27, 0xd2, 0x60, 0x4b };
static const uint8_t SECP256R1_Gx[] =
{ 0x6b, 0x17, 0xd1, 0xf2, 0xe1, 0x2c, 0x42, 0x47, 0xf8, 0xbc, 0xe6, 0xe5, 0x63,
    0xa4, 0x40, 0xf2, 0x77, 0x03, 0x7d, 0x81, 0x2d, 0xeb, 0x33, 0xa0, 0xf4,
    0xa1, 0x39, 0x45, 0xd8, 0x98, 0xc2, 0x96 };
static const uint8_t SECP256R1_Gy[] =
{ 0x4f, 0xe3, 0x42, 0xe2, 0xfe, 0x1a, 0x7f, 0x9b, 0x8e, 0xe7, 0xeb, 0x4a, 0x7c,
    0x0f, 0x9e, 0x16, 0x2b, 0xce, 0x33, 0x57, 0x6b, 0x31, 0x5e, 0xce, 0xcb,
    0xb6, 0x40, 0x68, 0x37, 0xbf, 0x51, 0xf5 };
static const uint8_t SECP256R1_Order[] =
{ 0xff, 0xff, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00, 0xff, 0xff, 0xff, 0xff, 0xff,
    0xff, 0xff, 0xff, 0xbc, 0xe6, 0xfa, 0xad, 0xa7, 0x17, 0x9e, 0x84, 0xf3,
    0xb9, 0xca, 0xc2, 0xfc, 0x63, 0x25, 0x51 };
static const uint8_t SECP256R1_OrderMinus1[] =
{ 0xff, 0xff, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00, 0xff, 0xff, 0xff, 0xff, 0xff,
    0xff, 0xff, 0xff, 0xbc, 0xe6, 0xfa, 0xad, 0xa7, 0x17, 0x9e, 0x84, 0xf3,
    0xb9, 0xca, 0xc2, 0xfc, 0x63, 0x25, 0x50 };

#ifdef ALIRO_TEST
// This is the static private key from the static test vectors
const uint8_t staticPrivateKey[] = {0x7a, 0x9e, 0x50, 0xa1, 0x9a, 0xe3, 0x85, 0xe3, 0x9b, 0x3b, 0xf0, 0xc7, 0x5e, 0xb5, 0xf9, 0xc9, 0xa5, 0xeb, 0x4d, 0x51, 0xf8, 0x8, 0x23, 0x18, 0x35, 0x39, 0x5f, 0xd2, 0xc1, 0x7, 0x83, 0x67};
#define STATIC_PRIVATE_KEY      staticPrivateKey
#else // ALIRO_TEST
// This is the static private key for the mobile phone app.
//ST OFFICE
//const uint8_t staticPrivateKey[] = {0xCB, 0x64, 0xA5, 0x8A, 0x5F, 0xAE, 0x05, 0x2D, 0x45, 0x8B, 0x3B, 0xCE, 0x25, 0x99, 0xA2, 0x0A,
//                                      0x66, 0x59, 0x87, 0xB5, 0x37, 0xB1, 0x57, 0xCC, 0x3C, 0xC4, 0x2E, 0x1D, 0x9C, 0x65, 0x24, 0x09};
//HOME
//const uint8_t staticPrivateKey[] = {0xBE, 0x08, 0x3F, 0xCD, 0xA3, 0x7B, 0xB2, 0xDA, 0x82, 0x66, 0xDC, 0x64, 0x61, 0x66, 0x35, 0x24,
//                                      0xE5, 0x9A, 0xD4, 0xA0, 0xBB, 0x6B, 0x4D, 0xD3, 0x1C, 0x88, 0x53, 0xC2, 0x0F, 0xFB, 0xB4, 0x0D};

#include "provisioning.h"

#define STATIC_PRIVATE_KEY      (getProvisioningFromFlash()->staticPrivateKey)
#endif // ALIRO_TEST

uint8_t ephemeralPrivateKey[] = {0x3c, 0xf, 0x74, 0x11, 0x4c, 0xd2, 0xa0, 0x21, 0xe8, 0x6, 0x6e, 0xfb, 0xaa, 0x31, 0xdb, 0xb9, 0x7e, 0xf0, 0x5, 0x42, 0x72, 0x19, 0x26, 0x6, 0xfd, 0x96, 0x63, 0x3a, 0x4, 0xf6, 0x62, 0x14};

#endif // USE_MCU_HW_ASYM

#ifndef GET_UINT32_BE
#define GET_UINT32_BE(n,b,i)                            \
    do {                                                    \
      (n) = ( (uint32_t) (b)[(i)    ] << 24 )             \
      | ( (uint32_t) (b)[(i) + 1] << 16 )             \
      | ( (uint32_t) (b)[(i) + 2] <<  8 )             \
      | ( (uint32_t) (b)[(i) + 3]       );            \
    } while( 0 )
#endif //GET_UINT32_BE


/**
 * @brief Initializes MCU cryptographic hardware modules.
 * @param rng Pointer to RNG handle.
 * @param pka Pointer to PKA handle.
 * @param hash Pointer to HASH handle.
 * @param cryp Pointer to CRYP handle.
 * @return SM_Crypto_Error_t error code.
 */
SM_Crypto_Error_t SM_initCryptoMCU(RNG_HandleTypeDef *rng, PKA_HandleTypeDef *pka, HASH_HandleTypeDef *hash, CRYP_HandleTypeDef *cryp)
{
#ifdef USE_MCU_HW_RNG
  HRng = rng;
#endif // USE_MCU_HW_RNG
#ifdef USE_MCU_HW_ASYM
  HPka = pka;
#endif // USE_MCU_HW_ASYM
  HHash = hash;
  HCryp = cryp;
  return SM_Crypto_OK;
}

#ifdef USE_MCU_HW_RNG

/**
 * @brief Generates random bytes using MCU hardware RNG.
 * @param bytes Pointer to output buffer.
 * @param size Number of bytes to generate.
 * @return SM_Crypto_Error_t error code.
 */
SM_Crypto_Error_t SM_genRandomBytes(uint8_t *bytes, uint32_t size)
{
  uint32_t temp;
  uint32_t remain = size % 4;
  if (HRng == NULL)
    return(SM_Crypto_HW_Error);
  __HAL_RCC_RNG_CLK_ENABLE();
  __HAL_RNG_ENABLE(HRng);
  
#ifdef STM32WBA65xx
  HW_RNG_Start();
#endif //STM32WBA65xx
  
  for (uint32_t i = 0; i < size / 4; i++)
  {
#ifdef STM32WBA65xx
	HW_RNG_Get(1, &temp);
#endif //STM32WBA65xx
	
#if defined(STM32U545xx) || defined(STM32U585xx) || defined(STM32U385xx)
    if (HAL_RNG_GenerateRandomNumber(HRng, &temp) != HAL_OK)
    {
      /* Random number generation error */
      return (SM_Crypto_HW_Error);
    }
#endif //defined(STM32U545xx) || defined(STM32U585xx) || defined(STM32U385xx)
#if defined(STM32WBA52xx)||defined(STM32WBA55xx)
    if (HAL_RNG_GenerateRandomNumber(HRng, &temp) != HAL_OK)
    {
      /* Random number generation error */
      return (SM_Crypto_HW_Error);
    }
#endif //defined(STM32WBA52xx)||defined(STM32WBA55xx)
    *(uint32_t *) bytes = temp;
    bytes += 4;
  }
  if (remain)
  {
#ifdef STM32WBA65xx
	HW_RNG_Get(1, &temp);
#endif //STM32WBA65xx

#if defined(STM32WBA52xx)||defined(STM32WBA55xx)
    if (HAL_RNG_GenerateRandomNumber(HRng, &temp) != HAL_OK)
    {
      /* Random number generation error */
      return (SM_Crypto_HW_Error);
    }
#endif //defined(STM32WBA52xx)||defined(STM32WBA55xx)
#if defined(STM32U545xx) || defined(STM32U585xx) || defined(STM32U385xx)
    if (HAL_RNG_GenerateRandomNumber(HRng, &temp) != HAL_OK)
    {
      /* Random number generation error */
      return (SM_Crypto_HW_Error);
    }
#endif //defined(STM32U545xx) || defined(STM32U585xx) || defined(STM32U385xx)
    for (uint32_t i = 0; i < remain; i++)
    {
      *bytes = (temp >> i) & 0xFF;
      bytes++;
    }
  }

  return SM_Crypto_OK;
}
#else //USE_MCU_HW_RNG
#if 0
SM_Crypto_Error_t SM_genRandomBytes(uint8_t *bytes, uint32_t size)
{
  for(uint32_t i = 0; i< size;i++)
  {
    *bytes = i;
  }
  return SM_Crypto_OK;
}
#endif
#endif // USE_MCU_HW_RNG

#ifdef USE_MCU_HW_ASYM
static uint32_t NumberOfWords(uint32_t NumberOfBytes)
{
  return ((NumberOfBytes + 3) / 4);
}
static SM_Crypto_Error_t GeneratePrivateKey(uint8_t *pPrivateKey)
{
  /* for NIST compliance (to decrease bias below 2^-64)
   * generate random with 64 extra bits (2 words)
   * perform modular reduction by n-1
   * and add 1 */
  uint8_t temp[MODULUS_SIZE_BYTES + 8] =   { 0 };
  uint8_t temp2[MODULUS_SIZE_BYTES + 8] =   { 0 };
  uint32_t random_bitstring_len_in_bits = 0;
  PKA_ModRedInTypeDef in_modred =   { 0 };

  random_bitstring_len_in_bits = SECP256R1_ORDER_LEN_BITS + 64;
  if (SM_genRandomBytes(temp, (random_bitstring_len_in_bits + 7) / 8) != SM_Crypto_OK)
  {
    return (SM_Crypto_HW_Error);
  }
  /* Perform modular reduction by n-1 */
  in_modred.pMod = SECP256R1_OrderMinus1;
  in_modred.modSize = sizeof(SECP256R1_Order);
  in_modred.pOp1 = (uint32_t*)temp;
  in_modred.OpSize = (random_bitstring_len_in_bits + 31) / 32;
  if (HAL_PKA_ModRed(HPka, &in_modred, PKA_TIMEOUT) != HAL_OK)
  {
    return (SM_Crypto_HW_Error);
  }
  HAL_PKA_Arithmetic_GetResult(HPka, (uint32_t*)temp2);

  /* compute reduced_random + 1 */
  memset(temp, 0, sizeof(temp));
  temp[0] = 1;
  {
    PKA_AddInTypeDef in_add;
    in_add.pOp1 = (uint32_t*)temp2;
    in_add.pOp2 = (uint32_t*)temp;
    in_add.size = NumberOfWords(sizeof(SECP256R1_Order));
    if (HAL_PKA_Add(HPka, &in_add, PKA_TIMEOUT) != HAL_OK)
    {
      return SM_Crypto_HW_Error;
    }
    HAL_PKA_Arithmetic_GetResult(HPka, (uint32_t*)temp2);
  }
  memcpy(pPrivateKey, temp2, sizeof(SECP256R1_Order));
  return SM_Crypto_OK;
}

static SM_Crypto_Error_t GeneratePublicKey(const uint8_t *pPrivateKey, uint8_t *pPublicKey)
{
  PKA_ECCMulInTypeDef in_mult =
  { 0 };

  /* Compute public key */
  in_mult.modulusSize = sizeof(SECP256R1_Prime);
  in_mult.coefSign = SECP256R1_A_SIGN;
  in_mult.coefA = SECP256R1_AbsA;
  in_mult.coefB = SECP256R1_B;
  in_mult.modulus = SECP256R1_Prime;
  in_mult.pointX = SECP256R1_Gx;
  in_mult.pointY = SECP256R1_Gy;
  in_mult.primeOrder = SECP256R1_Order;
  in_mult.scalarMulSize = sizeof(SECP256R1_Order); // WRONG NAMING IN HAL STRUCT: this must be the length of the ORDER
  in_mult.scalarMul = (uint8_t *)pPrivateKey;

  {
    PKA_ECCMulOutTypeDef mulRes;
    mulRes.ptX = &pPublicKey[0];
    mulRes.ptY = &pPublicKey[sizeof(SECP256R1_Order)];
    if (HAL_PKA_ECCMul(HPka, &in_mult, PKA_TIMEOUT) != HAL_OK)
    {
      return SM_Crypto_HW_Error;
    }
    HAL_PKA_ECCMul_GetResult(HPka, &mulRes);
  }
  return SM_Crypto_OK;
}


/**
 * @brief Performs ECDH key agreement using MCU hardware.
 * @param pPublicKey Pointer to public key.
 * @param pSharedSecret Pointer to output shared secret.
 * @return SM_Crypto_Error_t error code.
 */
SM_Crypto_Error_t SM_Ecdh(const uint8_t *pPublicKey, uint8_t *pSharedSecret)
{
  PKA_ECCMulInTypeDef in_mult =  { 0 };

  /* Compute ECDH */
  in_mult.modulusSize = sizeof(SECP256R1_Prime);
  in_mult.coefSign = SECP256R1_A_SIGN;
  in_mult.coefA = SECP256R1_AbsA;
  in_mult.coefB = SECP256R1_B;
  in_mult.modulus = SECP256R1_Prime;
  in_mult.pointX = (uint8_t*)&pPublicKey[0];
  in_mult.pointY = (uint8_t*)&pPublicKey[sizeof(SECP256R1_Order)];
  in_mult.primeOrder = SECP256R1_Order;
  in_mult.scalarMulSize = sizeof(SECP256R1_Order); // WRONG NAMING IN HAL STRUCT: this must be the length of the ORDER
  in_mult.scalarMul = ephemeralPrivateKey;

  if (HAL_PKA_ECCMul(HPka, &in_mult, PKA_TIMEOUT) != HAL_OK)
  {
    memset(ephemeralPrivateKey, 0, sizeof(ephemeralPrivateKey));
    return SM_Crypto_Error;
  }
  memset(ephemeralPrivateKey, 0, sizeof(ephemeralPrivateKey));
  {
    PKA_ECCMulOutTypeDef pub;
    uint8_t tmp[sizeof(SECP256R1_Order)];
    pub.ptX = pSharedSecret;
    pub.ptY = tmp;
    HAL_PKA_ECCMul_GetResult(HPka, &pub);
  }
  return SM_Crypto_OK;
}


/**
 * @brief Signs a digest using MCU hardware.
 * @param digest Pointer to digest to sign.
 * @param digestSize Size of digest.
 * @param signature Pointer to output signature.
 * @return SM_Crypto_Error_t error code.
 */
SM_Crypto_Error_t SM_sign(const uint8_t *digest, uint32_t digestSize, uint8_t *signature)
{
  uint8_t a_k[MODULUS_SIZE_WORDS*sizeof(uint32_t)] =
  { 0 };
  PKA_ECDSASignInTypeDef in =
  { 0 };

  if (digestSize != SHA256_BYTE_SIZE)
    return SM_Crypto_Error;
  /* generate random k parameter in NIST compliant way */
  if (GeneratePrivateKey(a_k) != SM_Crypto_OK)
  {
    return (SM_Crypto_HW_Error);
  }

  /* Set input parameters */
  in.primeOrderSize = sizeof(SECP256R1_Order);
  in.modulusSize = sizeof(SECP256R1_Prime);
  in.coefSign = SECP256R1_A_SIGN;
  in.coef = SECP256R1_AbsA;
  in.coefB = SECP256R1_B;
  in.modulus = SECP256R1_Prime;
  in.basePointX = SECP256R1_Gx;
  in.basePointY = SECP256R1_Gy;
  in.primeOrder = SECP256R1_Order;

  in.integer = a_k;
  in.hash = digest;
  
  in.privateKey = STATIC_PRIVATE_KEY;

  if (HAL_PKA_ECDSASign(HPka, &in, PKA_TIMEOUT) != HAL_OK)
  {
    return (SM_Crypto_HW_Error);
  }
  {
    PKA_ECDSASignOutTypeDef out;
    out.RSign = &signature[0];
    out.SSign = &signature[sizeof(SECP256R1_Order)];
    HAL_PKA_ECDSASign_GetResult(HPka, &out, NULL);
  }
  return SM_Crypto_OK;
}


/**
 * @brief Generates a keypair using MCU hardware.
 * @param pub Pointer to output public key.
 * @return SM_Crypto_Error_t error code.
 */
SM_Crypto_Error_t SM_genKeypair(uint8_t *pub)
{
  if (GeneratePrivateKey(ephemeralPrivateKey) != SM_Crypto_OK)
    return SM_Crypto_Error;

  if (GeneratePublicKey(ephemeralPrivateKey, pub) != SM_Crypto_OK)
  {
    memset(ephemeralPrivateKey, 0, sizeof(ephemeralPrivateKey));
    return SM_Crypto_Error;
  }
  return SM_Crypto_OK;
}

/**
 * @brief Verifies a signature using MCU hardware.
 * @param pPublicKey Pointer to public key.
 * @param digest Pointer to digest.
 * @param digestSize Size of digest.
 * @param signature Pointer to signature.
 * @return SM_Crypto_Error_t error code.
 */
SM_Crypto_Error_t SM_verify(const uint8_t *pPublicKey, const uint8_t *digest, const uint32_t digestSize, const uint8_t *signature)
{
  PKA_ECDSAVerifInTypeDef in =
  { 0 };
  if (digestSize != SHA256_BYTE_SIZE)
    return (SM_Crypto_Error);
  /* Set input parameters */
  in.primeOrderSize = sizeof(SECP256R1_Order);
  in.modulusSize = sizeof(SECP256R1_Prime);
  in.coefSign = SECP256R1_A_SIGN;
  in.coef = SECP256R1_AbsA;
  in.modulus = SECP256R1_Prime;
  in.basePointX = SECP256R1_Gx;
  in.basePointY = SECP256R1_Gy;
  in.primeOrder = SECP256R1_Order;

  in.RSign = &signature[0];
  in.SSign = &signature[PRIVATE_KEY_SIZE_BYTES];
  in.hash = digest;
  in.pPubKeyCurvePtX = &pPublicKey[0];
  in.pPubKeyCurvePtY = &pPublicKey[MODULUS_SIZE_BYTES];

  if (HAL_PKA_ECDSAVerif(HPka, &in, PKA_TIMEOUT) != HAL_OK)
  {
    return (SM_Crypto_HW_Error);
  }
  if (!HAL_PKA_ECDSAVerif_IsValidSignature(HPka))
    return (SM_Crypto_Error);

  return SM_Crypto_OK;
}


/**
 * @brief Gets the certificate from MCU hardware.
 * @param certificate Pointer to output certificate buffer.
 * @param certificateSize Pointer to output certificate size.
 * @return SM_Crypto_Error_t error code.
 */
SM_Crypto_Error_t SM_getCertificate(uint8_t *certificate, uint32_t *certificateSize)
{
  return SM_Crypto_OK;
}

#endif // USE_MCU_HW_ASYM









/* Private function -----------------------------------------------*/
static SM_Crypto_Error_t KDFKeyExpansionRound(const uint8_t* pKDK, const uint8_t *prev, const uint8_t* pFixedInfo, uint32_t fixedInfoSize, uint8_t counter, uint8_t* output)
{
  /* counter and L must be big-endian, update least significant byte */
  uint32_t offset = 0;
  uint32_t hmac_input_size = 0;

  uint8_t input[SHA256_BYTE_SIZE + MAX_FIXED_INFO_SIZE + 1 ];

  if (fixedInfoSize > MAX_FIXED_INFO_SIZE)
    return SM_Crypto_Error;
  HHash->Init.DataType = HHASH_INIT_DATATYPE;
  HASH_SET_ALGORITHM(HHash->Init, HASH_ALGOSELECTION_SHA256);
  HHash->Init.KeySize = SHA256_BYTE_SIZE;
  HHash->Init.pKey = (uint8_t *)pKDK;
  if (HAL_HASH_Init(HHash) != HAL_OK)
  {
    return SM_Crypto_HW_Error;
  }
  hmac_input_size = SHA256_BYTE_SIZE + fixedInfoSize + 1 ;
  if (counter != 1)
  {
    memcpy(input + offset, prev, SHA256_BYTE_SIZE);
    offset += SHA256_BYTE_SIZE;
  } else {
    /* No need to copy, but will hmac 32 bytes less */
    hmac_input_size -= SHA256_BYTE_SIZE;
  }
  /* Based on RFC 5869 */
  memcpy(input + offset, pFixedInfo, fixedInfoSize);
  offset += fixedInfoSize;
  memcpy(input + offset, &counter, 1);
  offset += 1;
  if (SM_HMAC_SHA256_START(HHash, input, hmac_input_size, output, 5000) != HAL_OK)
  {
    return SM_Crypto_HW_Error;
  }
  return SM_Crypto_OK;
}



/* Public functions -----------------------------------------------*/

/**
 * @brief Computes SHA-1 hash using MCU hardware.
 * @param input Pointer to input data.
 * @param inputSize Size of input data.
 * @param hash Pointer to output hash buffer.
 * @return SM_Crypto_Error_t error code.
 */
SM_Crypto_Error_t SM_sha1(uint8_t *input, uint32_t inputSize, uint8_t *hash)
{
  HHash->Init.DataType = HHASH_INIT_DATATYPE;
  HASH_SET_ALGORITHM(HHash->Init, HASH_ALGOSELECTION_SHA1);
  if (HAL_HASH_Init(HHash) != HAL_OK)
  {
    return SM_Crypto_HW_Error;
  }

  if (SM_HASH_SHA1_START(HHash, input, inputSize, &hash[0], HASH_TIMEOUT) != HAL_OK)
  {
    return (SM_Crypto_HW_Error);
  }
  return SM_Crypto_OK;
}



/**
 * @brief Computes SHA-256 hash using MCU hardware.
 * @param input Pointer to input data.
 * @param inputSize Size of input data.
 * @param hash Pointer to output hash buffer.
 * @return SM_Crypto_Error_t error code.
 */
SM_Crypto_Error_t SM_sha256(uint8_t *input, uint32_t inputSize, uint8_t *hash)
{
  HHash->Init.DataType = HHASH_INIT_DATATYPE;
  HASH_SET_ALGORITHM(HHash->Init, HASH_ALGOSELECTION_SHA256);
  if (HAL_HASH_Init(HHash) != HAL_OK)
  {
    return SM_Crypto_HW_Error;
  }

  if (SM_HASH_SHA256_START(HHash, input, inputSize, &hash[0], HASH_TIMEOUT) != HAL_OK)
  {
    return (SM_Crypto_HW_Error);
  }
  return SM_Crypto_OK;
}


/**
 * @brief Starts SHA-256 computation using MCU hardware.
 * @return SM_Crypto_Error_t error code.
 */
SM_Crypto_Error_t SM_sha256_start(void)
{
  HHash->Init.DataType = HHASH_INIT_DATATYPE;
  HASH_SET_ALGORITHM(HHash->Init, HASH_ALGOSELECTION_SHA256);
  if (HAL_HASH_Init(HHash) != HAL_OK)
  {
    return SM_Crypto_HW_Error;
  }
  return SM_Crypto_OK;
}


/**
 * @brief Updates SHA-256 computation with new data using MCU hardware.
 * @param input Pointer to input data.
 * @param inputSize Size of input data.
 * @return SM_Crypto_Error_t error code.
 */
SM_Crypto_Error_t SM_sha256_update(const uint8_t *input, uint32_t inputSize)
{
  if (SM_HASH_SHA256_ACCUMULATE(HHash, input, inputSize, HASH_TIMEOUT) != HAL_OK)
  {
    return (SM_Crypto_HW_Error);
  }
  return SM_Crypto_OK;
}


/**
 * @brief Finishes SHA-256 computation and returns the hash using MCU hardware.
 * @param hash Pointer to output hash buffer.
 * @param input Pointer to input data.
 * @param inputSize Size of input data.
 * @return SM_Crypto_Error_t error code.
 */
SM_Crypto_Error_t SM_sha256_finish(uint8_t *hash, const uint8_t *input, uint32_t inputSize)
{
  if (SM_HASH_SHA256_ACCUMULATE_LAST(HHash, input, inputSize, &hash[0], HASH_TIMEOUT) != HAL_OK)
  {
    return (SM_Crypto_HW_Error);
  }
  return SM_Crypto_OK;
}


/**
 * @brief Performs randomness extraction for KDF.
 * @param Z Pointer to input keying material.
 * @param ZSize Size of input keying material.
 * @param salt Pointer to salt.
 * @param saltSize Size of salt.
 * @param output Pointer to output buffer.
 * @return SM_Crypto_Error_t error code.
 */
SM_Crypto_Error_t KDFRandomnessExtraction(const uint8_t* Z, uint32_t ZSize,
    const uint8_t* salt, uint32_t saltSize,
    uint8_t* output)
{
  uint8_t null_salt[SHA256_BYTE_SIZE] = {0};
  
  if ((saltSize == 0) && (salt == NULL)) // no salt (optionnal), need to use string of zeros, the size is the size of the HASH
  {
    HHash->Init.KeySize = SHA256_BYTE_SIZE;
    HHash->Init.pKey = null_salt;
  } else if (salt == NULL) // saltsize is non-zero but pointer is null -> error
  {
    return SM_Crypto_Error;
  }
  else
  {
    HHash->Init.KeySize = saltSize;
    HHash->Init.pKey = (uint8_t *)salt;
  }
  HHash->Init.DataType = HHASH_INIT_DATATYPE;
  HASH_SET_ALGORITHM(HHash->Init, HASH_ALGOSELECTION_SHA256);
  if (HAL_HASH_Init(HHash) != HAL_OK)
  {
    return SM_Crypto_HW_Error;
  }
  if (SM_HMAC_SHA256_START(HHash, (uint8_t *)Z, ZSize, output, 5000) != HAL_OK)
  {
    return SM_Crypto_HW_Error;
  }

  return SM_Crypto_OK;
}


/**
 * @brief Key Derivation Function RFC5869 using MCU hardware.
 * @param pZ Pointer to input keying material.
 * @param ZSize Size of input keying material.
 * @param pSalt Pointer to salt.
 * @param saltSize Size of salt.
 * @param pFixedInfo Pointer to fixed info.
 * @param fixedInfoSize Size of fixed info.
 * @param L Length of output keying material.
 * @param pOutput Pointer to output buffer.
 * @return SM_Crypto_Error_t error code.
 */
SM_Crypto_Error_t SM_KDFRFC5869(const uint8_t* pZ, uint32_t ZSize,
    const uint8_t* pSalt, uint32_t saltSize,
    const uint8_t* pFixedInfo, uint32_t fixedInfoSize,
    uint32_t L, uint8_t* pOutput)
{
  uint8_t counter = 0;
  uint8_t KDK[SHA256_BYTE_SIZE] = {0};
  uint8_t *prev = NULL;

  if(L >= 255*SHA256_BYTE_SIZE)
  {
    return SM_Crypto_HW_Error;
  }
  if (KDFRandomnessExtraction(pZ, ZSize, pSalt, saltSize, KDK) != SM_Crypto_OK)
  {
    return SM_Crypto_HW_Error;
  }
  for(counter=1; counter<=L/SHA256_BYTE_SIZE; counter++)
  {
    /* compute full rounds */
    if(KDFKeyExpansionRound(KDK, prev, pFixedInfo, fixedInfoSize, counter, pOutput) != SM_Crypto_OK)
    {
      return SM_Crypto_HW_Error;
    }
    prev = pOutput;
    pOutput += SHA256_BYTE_SIZE;
  }
  if(L % SHA256_BYTE_SIZE)
  {
    /* compute non full round */
    uint8_t buffer[SHA256_BYTE_SIZE];
    if(KDFKeyExpansionRound(KDK, prev, pFixedInfo, fixedInfoSize, counter, buffer) != SM_Crypto_OK)
    {
      return SM_Crypto_HW_Error;
    }
    memcpy(pOutput, buffer, L % SHA256_BYTE_SIZE);
  }
  return SM_Crypto_OK;
}


/**
 * @brief Derives KDH key using MCU hardware.
 * @param pZ Pointer to input keying material.
 * @param info Pointer to info data.
 * @param infoSize Size of info data.
 * @param pKdh Pointer to output KDH key.
 * @return SM_Crypto_Error_t error code.
 */
SM_Crypto_Error_t SM_deriveKdh(const uint8_t *pZ, const uint8_t * info, uint32_t infoSize, uint8_t *pKdh)
{
  uint8_t counter_be[4] = {0,0,0,1};
  uint32_t offset = 0;
  uint8_t buf[KDH_INFO_MAX_LEN_BYTES + 32 + 4];
  memcpy(&buf[offset], pZ, 32);
  offset += 32;
  memcpy(&buf[offset], counter_be, 4);
  offset += 4;
  memcpy(&buf[offset], info, infoSize);

  HHash->Init.DataType = HHASH_INIT_DATATYPE;
  HASH_SET_ALGORITHM(HHash->Init, HASH_ALGOSELECTION_SHA256);
  if (HAL_HASH_Init(HHash) != HAL_OK)
  {
    return SM_Crypto_HW_Error;
  }
  if (SM_HASH_SHA256_START(HHash, buf, 32 + 4 + infoSize, pKdh, 5000) != HAL_OK)
  {
    return SM_Crypto_HW_Error;
  }

  return SM_Crypto_OK;
}


/**
 * @brief Decrypts data using AES-256 GCM with MCU hardware.
 * @param pKey Pointer to key.
 * @param pNonce Pointer to nonce.
 * @param pAad Pointer to additional authenticated data.
 * @param aadSize Size of AAD.
 * @param pPayload Pointer to encrypted payload.
 * @param payloadSize Size of payload.
 * @param pPlaintext Pointer to output plaintext buffer.
 * @param pTag Pointer to authentication tag.
 * @return SM_Crypto_Error_t error code.
 */
SM_Crypto_Error_t SM_aes256GcmDecrypt(const uint8_t* pKey, const uint8_t* pNonce, const uint8_t* pAad, const uint32_t aadSize,
    const uint8_t* pPayload, const uint32_t payloadSize, uint8_t* pPlaintext,
    const uint8_t* pTag)
{
  uint32_t i = 0;
  uint32_t key[8] = {0};
  uint32_t iv[4] = {0};
  uint32_t tag[4] = {0};
  uint8_t tmp = 0;

  /* !!! Beware: the last word must be padded with 0s by the caller, otherwise tag will be wrong !!! */
  if (payloadSize % 4)
  {
    /* Must check that the padding has been properly done by caller */
    if (memcmp(&pPayload[payloadSize], &i, 4 - (payloadSize % 4)))
      return SM_Crypto_Error;
  }
  /* Format and fill AES key  */
  for( i=0; i < 8; i++ )
    GET_UINT32_BE( key[i], pKey, 4*i );
  /* Set IV with invert endianness */
  for( i=0; i < 3; i++ )
    GET_UINT32_BE( iv[i], pNonce, 4*i );

  /* counter value must be set to 2 when processing the first block of payload */
  iv[3] = 0x00000002;

  HCryp->Instance = AES;
  HCryp->Init.DataType = CRYP_DATATYPE_8B;
  HCryp->Init.KeySize = CRYP_KEYSIZE_256B;
  HCryp->Init.pKey = key;
  HCryp->Init.pInitVect = iv;
  HCryp->Init.Algorithm = CRYP_AES_GCM_GMAC;
  HCryp->Init.Header = (uint32_t *)pAad;
  HCryp->Init.HeaderSize = aadSize;
  HCryp->Init.DataWidthUnit = CRYP_DATAWIDTHUNIT_BYTE;
  HCryp->Init.HeaderWidthUnit = CRYP_HEADERWIDTHUNIT_BYTE;
  HCryp->Init.KeyIVConfigSkip = CRYP_KEYIVCONFIG_ALWAYS;
  HCryp->Init.KeyMode = CRYP_KEYMODE_NORMAL;
  if (HAL_CRYP_Init(HCryp) != HAL_OK)
  {
    return(SM_Crypto_HW_Error);
  }
  if (HAL_CRYP_Decrypt(HCryp, (uint32_t *)pPayload, (uint16_t) payloadSize, (uint32_t *)pPlaintext, 5000) != HAL_OK)
    return(SM_Crypto_HW_Error);
  /* Last 16 bytes of plaintext is the GCM tag */
  if (HAL_CRYPEx_AESGCM_GenerateAuthTAG(HCryp, (const uint32_t *)tag, 5000) != HAL_OK)
  {
    return SM_Crypto_HW_Error;
  }
  // constant time compare
  for (i = 0; i < 4; i++)
  {
    tmp |= pTag[i*4]     ^ (uint8_t)(tag[i] & 0xFF);
    tmp |= pTag[i*4 + 1] ^ (uint8_t)((tag[i] >> 8) & 0xFF);
    tmp |= pTag[i*4 + 2] ^ (uint8_t)((tag[i] >> 16) & 0xFF);
    tmp |= pTag[i*4 + 3] ^ (uint8_t)((tag[i] >> 24) & 0xFF);
  }
  if (tmp)
  {
    memset(pPlaintext, 0, payloadSize);
    return SM_Crypto_Error;
  }

  return SM_Crypto_OK;
}


/**
 * @brief Encrypts data using AES-256 GCM with MCU hardware.
 * @param pKey Pointer to key.
 * @param pNonce Pointer to nonce.
 * @param pAad Pointer to additional authenticated data.
 * @param aadSize Size of AAD.
 * @param pInput Pointer to input data.
 * @param inputSize Size of input data.
 * @param pOutput Pointer to output buffer.
 * @param pOutputSize Pointer to output size.
 * @return SM_Crypto_Error_t error code.
 */
SM_Crypto_Error_t SM_aes256GcmEncrypt(const uint8_t* pKey, const uint8_t* pNonce, const uint8_t* pAad, const uint32_t aadSize,
    const uint8_t* pInput, const uint32_t inputSize, uint8_t* pOutput, uint32_t* pOutputSize)
{
  uint32_t i = 0;
  uint32_t key[8] = {0};
  uint32_t iv[4] = {0};
  uint32_t tag[4] = {0};
 
  
  /* Workaround HAL: need to pad with 0s the last word of payload. Here we first copy the input into the output
     buffer (that is at least 16 bytes longer (for the tag), and then we encrypt in place */
  memcpy(pOutput, pInput, inputSize);
  if ((inputSize) % 4)
    memset(&pOutput[inputSize], 0, 4 - ((inputSize) % 4));
  
  /* Format and fill AES key  */
  for( i=0; i < 8; i++ )
    GET_UINT32_BE( key[i], pKey, 4*i );
  /* Set IV with invert endianness */
  for( i=0; i < 3; i++ )
    GET_UINT32_BE( iv[i], pNonce, 4*i );

  /* counter value must be set to 2 when processing the first block of payload */
  iv[3] = 0x00000002;

  HCryp->Instance = AES;
  HCryp->Init.DataType = CRYP_DATATYPE_8B;
  HCryp->Init.KeySize = CRYP_KEYSIZE_256B;
  HCryp->Init.pKey = key;
  HCryp->Init.pInitVect = iv;
  HCryp->Init.Algorithm = CRYP_AES_GCM_GMAC;
  HCryp->Init.Header = (uint32_t *)pAad;
  HCryp->Init.HeaderSize = aadSize;
  HCryp->Init.DataWidthUnit = CRYP_DATAWIDTHUNIT_BYTE;
  HCryp->Init.HeaderWidthUnit = CRYP_HEADERWIDTHUNIT_BYTE;
  HCryp->Init.KeyIVConfigSkip = CRYP_KEYIVCONFIG_ALWAYS;
  HCryp->Init.KeyMode = CRYP_KEYMODE_NORMAL;
  if (HAL_CRYP_Init(HCryp) != HAL_OK)
  {
    return(SM_Crypto_HW_Error);
  }
  if (HAL_CRYP_Encrypt(HCryp, (uint32_t *)pOutput, (uint16_t)inputSize, (uint32_t *)pOutput, 5000) != HAL_OK)
    return(SM_Crypto_HW_Error);
  /* Generate the GCM tag */
  if (HAL_CRYPEx_AESGCM_GenerateAuthTAG(HCryp, (const uint32_t *)tag, 5000) != HAL_OK)
  {
    return SM_Crypto_HW_Error;
  }
  for (i = 0; i < 4; i++) {
    pOutput[inputSize + i*4] = (uint8_t)(tag[i] & 0xFF);
    pOutput[inputSize + i*4 + 1] = (uint8_t)((tag[i] >> 8) & 0xFF);
    pOutput[inputSize + i*4 + 2] = (uint8_t)((tag[i] >> 16) & 0xFF);
    pOutput[inputSize + i*4 + 3] = (uint8_t)((tag[i] >> 24) & 0xFF);
  }
  *pOutputSize = inputSize + 16; 
  return SM_Crypto_OK;
}

