/*
* Creation date : May 29 15:33:10 2007
* Last modified : %modify_time%
*/
/** @file
* \brief This file contains implementation of 
* LLF_KDF() function. 
*
* \version LLF_KDF.c#1:csrc:1
* \author Yermalayeu Ihar
* \remarks Copyright (C) 2007 by Discretix Technologies Ltd.
* All Rights reserved
*/

/************************ Include Files ***********************/

#include "CE2_KDF.h"
#include "LLF_KDF.h"
#include "tomcrypt.h"
#include <stdlib.h>
#include <memory.h>

/************************ Defines *****************************/

#define LLF_KDF_COUNTER_SIZE 4
#define LLF_KDF_MAX_OID_NUMBER 100
#define LLF_KDF_OTHER_INFO_SIZE 6

/************************ Enums *******************************/
/************************ Typedefs ****************************/
/************************ Global Data *************************/
/************************ Private function prototype **********/
/************************ Private Functions *******************/

#if 0
void PrintBuffer(FILE* out, DxUint8_t* buffer, DxUint32_t size, DxUint8_t* comment)
{
  DxUint32_t i;

  if (comment) 
    fprintf(out, "%s", comment);
  for(i = 0; i < size; i++) {
    if (i%16 == 0)
      fprintf(out, "\n");
    fprintf(out, " ");
    if (buffer[i] < 0x10)
      fprintf(out, "0");
    fprintf(out, "%X", buffer[i]);
  }
}
#endif


/**
****************************************************************
* Function Name: 
*  LLF_KDF_ASN1_KeyDerivFunc
*
*  @param[in] ZZSecret_ptr 	- A pointer to shared secret value octet string.
*  @param[in] ZZSecretSize  - The shared secret key Size, in bytes.
*  @param[in] OtherInfo     - The pointer to structure, containing pointers 
*                             and sizes of optional data shared by two
*                             entities intended to share the secret value. 
*                             This argument (structure) and also its members 
*                             are optional (if any is not need - set its pointer 
*                             and size to NULL).
*  @param[in] KDFhashMode	  - The hash function to be used. The hash function 
*                             output must be at least 160 bits.
*  @param[out] KeyingData_ptr - A pointer to the keying data derived from the 
*                               secret key, of length KeyLenInBits
*  @param[in] KeyLenInBytes	 - The size in bytes of the keying data to be generated. 
*                              In our implementation - KeyLenInBytes <= 2^32-1  
*                              (in standards KeyLengthInBits < hashlen*(2^32-1)).
*
* @returns \b
*  CE2Error_t 
*  - CE2_OK - on success
*  - Otherwise - on failure a value from CE2_error.h:
*
* \brief \b 
* Description:
*  LLF_KDF_ASN1_KeyDerivFunc performs key derivation according to ASN1 DER 
*  encoding method defined in standard  ANSI X9.42-2001, 7.2.1.
*  
*  \b 
* Algorithm:
*  -# Init variables 
*  -# Derivate key data 
***************************************************************/
CE2Error_t LLF_KDF_ASN1_KeyDerivFunc(DxUint8_t                *ZZSecret_ptr,
                                     DxUint32_t                ZZSecretSize,
                                     CE2_KDF_OtherInfo_t     *OtherInfo_ptr,
                                     CE2_KDF_HASH_OpMode_t    KDFhashMode,
                                     DxUint8_t                *KeyingData_ptr,
                                     DxUint32_t                KeyLenInBytes)
{
  CE2Error_t result = CE2_OK, error;
  DxUint32_t hashSize, d, i, i_counter, bufferSize, pos, counterPos;
  DxUint8_t *pHashes = DX_NULL, *pBuffer = DX_NULL;
  CE2_HASH_OperationMode_t hashMode;
  CE2_HASH_Result_t hashResult;
  
  /******************/
  /* Init variables */
  /******************/
  switch (KDFhashMode) {
    case CE2_KDF_HASH_SHA1_mode:
      hashSize = CE2_HASH_SHA1_DIGEST_SIZE_IN_BYTES;
      hashMode = CE2_HASH_SHA1_mode;
      break;
    case CE2_KDF_HASH_SHA224_mode:
      hashSize = CE2_HASH_SHA224_DIGEST_SIZE_IN_BYTES;
      hashMode = CE2_HASH_SHA224_mode;
      break;
    case CE2_KDF_HASH_SHA256_mode:
      hashSize = CE2_HASH_SHA256_DIGEST_SIZE_IN_BYTES;
      hashMode = CE2_HASH_SHA256_mode;
      break;
    case CE2_KDF_HASH_SHA384_mode:
      hashSize = CE2_HASH_SHA384_DIGEST_SIZE_IN_BYTES;
      hashMode = CE2_HASH_SHA384_mode;
      break;
    case CE2_KDF_HASH_SHA512_mode:
      hashSize = CE2_HASH_SHA512_DIGEST_SIZE_IN_BYTES;
      hashMode = CE2_HASH_SHA512_mode;
      break;
    default:
      return CE2_LLF_KDF_MODULE_ERROR_BASE;
  }

  /* d = [keyLen/hashLen] */
  d = (KeyLenInBytes + hashSize -1)/hashSize + 1;

  /* Allocate hash buffer */
  pHashes = malloc(d*hashSize*sizeof(DxUint8_t));
  if (pHashes == DX_NULL) {
    result = CE2_LLF_KDF_MODULE_ERROR_BASE;
    goto error_case;
  }

  /***********************************************/                             
  /* counter = 1                                 */
  /* For i = 1 to d                              */
  /*   Hash_i = H(ZZ||OtherInfo);                */
  /*   OtherInfo ::= AlgorithmID || Counter ||   */
  /*    Optional(PartyUInfo || PartyVInfo ||     */
  /*    SuppPrivInfo || SuppPubInfo)             */
  /*   Counter = Counter + 1                     */ 
  /*   Increment i                               */
  /***********************************************/                             

  /* Determinate size and allocate buffer for ZZ||OtherInfo */
  bufferSize = ZZSecretSize + LLF_KDF_COUNTER_SIZE;
  bufferSize += OtherInfo_ptr->SizeOfAlgorithmID;
  if (OtherInfo_ptr->PartyUInfo_ptr != DX_NULL)
    bufferSize += OtherInfo_ptr->SizeOfPartyUInfo;
  if (OtherInfo_ptr->PartyVInfo_ptr != DX_NULL) 
    bufferSize += OtherInfo_ptr->SizeOfPartyVInfo;
  if (OtherInfo_ptr->SuppPrivInfo_ptr != DX_NULL) 
    bufferSize += OtherInfo_ptr->SizeOfSuppPrivInfo;
  if (OtherInfo_ptr->SuppPubInfo_ptr != DX_NULL) 
    bufferSize += OtherInfo_ptr->SizeOfSuppPubInfo;
  pBuffer = malloc(bufferSize);
  if (pBuffer == DX_NULL) {
    result = CE2_LLF_KDF_MODULE_ERROR_BASE;
    goto error_case;
  }

  /* Set ZZSecret to buffer */ 
  pos = 0;
  memcpy(pBuffer + pos, ZZSecret_ptr, ZZSecretSize);
  pos += ZZSecretSize;

  /* Set AlgorithmID to buffer */ 
  memcpy(pBuffer + pos, OtherInfo_ptr->AlgorithmID_ptr, 
    OtherInfo_ptr->SizeOfAlgorithmID);
  pos += OtherInfo_ptr->SizeOfAlgorithmID;

  /* Leave empty space for Counter */ 
  counterPos = pos;
  pos += LLF_KDF_COUNTER_SIZE;

  /* Optional set PartyUInfo, PartyVInfo, SuppPubInfo, or SuppPrivInfo */ 
  if (OtherInfo_ptr->PartyUInfo_ptr != DX_NULL) {
    memcpy(pBuffer + pos, OtherInfo_ptr->PartyUInfo_ptr, 
      OtherInfo_ptr->SizeOfPartyUInfo);
    pos += OtherInfo_ptr->SizeOfPartyUInfo;
  }
  if (OtherInfo_ptr->PartyVInfo_ptr != DX_NULL) {
    memcpy(pBuffer + pos, OtherInfo_ptr->PartyVInfo_ptr, 
      OtherInfo_ptr->SizeOfPartyVInfo);
    pos += OtherInfo_ptr->SizeOfPartyVInfo;
  }
  if (OtherInfo_ptr->SuppPubInfo_ptr != DX_NULL) {
    memcpy(pBuffer + pos, OtherInfo_ptr->SuppPubInfo_ptr, 
      OtherInfo_ptr->SizeOfSuppPubInfo);
    pos += OtherInfo_ptr->SizeOfSuppPubInfo;
  }
  if (OtherInfo_ptr->SuppPrivInfo_ptr != DX_NULL) {
    memcpy(pBuffer + pos, OtherInfo_ptr->SuppPrivInfo_ptr, 
      OtherInfo_ptr->SizeOfSuppPrivInfo);
    pos += OtherInfo_ptr->SizeOfSuppPrivInfo;
  }

  i_counter = 1;
  for (i = 0; i < d; i++) {
    /* Write i_counter as big-endian 32-bit octed string Counter */ 
    pBuffer[counterPos + 0] = (DxUint8_t)(i_counter >> 24);
    pBuffer[counterPos + 1] = (DxUint8_t)(i_counter >> 16);
    pBuffer[counterPos + 2] = (DxUint8_t)(i_counter >> 8);
    pBuffer[counterPos + 3] = (DxUint8_t)(i_counter);

    /* Hash_i = H(ZZ||OtherInfo) */
    error = CE2_HASH(hashMode, pBuffer, bufferSize, hashResult);
    if (error != CE2_OK) {
      result = error;
      goto error_case;
    }

    memcpy(pHashes + i*hashSize, hashResult, hashSize);
    i_counter++;
  }

  /* KeyingData = leftmost keylen bits of Hash1||Hash2||…||Hashd */
  memcpy(KeyingData_ptr, pHashes, KeyLenInBytes);

error_case:
  if (pHashes != DX_NULL)
    free(pHashes);
  if (pBuffer != DX_NULL)
    free(pBuffer);
  return result;
} /* End of LLF_KDF_ASN1_KeyDerivFunc */

/**
****************************************************************
* Function Name: 
*  LLF_KDF_ConcatKeyDerivFunc
*
*  @param[in] ZZSecret_ptr 	- A pointer to shared secret value octet string.
*  @param[in] ZZSecretSize  - The shared secret key Size, in bytes.
*  @param[in] OtherInfo     - The pointer to structure, containing pointers 
*                             and sizes of optional data shared by two
*                             entities intended to share the secret value. 
*                             This argument (structure) and also its members 
*                             are optional (if any is not need - set its pointer 
*                             and size to NULL).
*  @param[in] KDFhashMode	  - The hash function to be used. The hash function 
*                             output must be at least 160 bits.
*  @param[out] KeyingData_ptr - A pointer to the keying data derived from the 
*                               secret key, of length KeyLenInBits
*  @param[in] KeyLenInBytes	 - The size in bytes of the keying data to be generated. 
*                              In our implementation - KeyLenInBytes <= 2^32-1  
*                              (in standards KeyLengthInBits < hashlen*(2^32-1)).
*
* @returns \b
*  CE2Error_t 
*  - CE2_OK - on success
*  - Otherwise - on failure a value from CE2_error.h:
*
* \brief \b 
* Description:
*  LLF_KDF_ConcatKeyDerivFunc performs key derivation according to concatenation 
*  mode defined in standard  ANSI X9.42-2001, 7.2.2.
*  
*  \b 
* Algorithm:
*  -# Init variables 
*  -# Derivate key data 
***************************************************************/

CE2Error_t LLF_KDF_ConcatKeyDerivFunc(DxUint8_t                *ZZSecret_ptr,
                                      DxUint32_t                ZZSecretSize,
                                      CE2_KDF_OtherInfo_t     *OtherInfo_ptr,
                                      CE2_KDF_HASH_OpMode_t    KDFhashMode,
                                      DxUint8_t                *KeyingData_ptr,
                                      DxUint32_t                KeyLenInBytes )
{
  CE2Error_t result = CE2_OK, error;
  DxUint32_t hashSize, d, i, i_counter, bufferSize, pos, counterPos;
  DxUint8_t *pHashes = DX_NULL, *pBuffer = DX_NULL;
  CE2_HASH_OperationMode_t hashMode;
  CE2_HASH_Result_t hashResult;

  /******************/
  /* Init variables */
  /******************/
  switch (KDFhashMode) {
    case CE2_KDF_HASH_SHA1_mode:
      hashSize = CE2_HASH_SHA1_DIGEST_SIZE_IN_BYTES;
      hashMode = CE2_HASH_SHA1_mode;
      break;
    case CE2_KDF_HASH_SHA224_mode:
      hashSize = CE2_HASH_SHA224_DIGEST_SIZE_IN_BYTES;
      hashMode = CE2_HASH_SHA224_mode;
      break;
    case CE2_KDF_HASH_SHA256_mode:
      hashSize = CE2_HASH_SHA256_DIGEST_SIZE_IN_BYTES;
      hashMode = CE2_HASH_SHA256_mode;
      break;
    case CE2_KDF_HASH_SHA384_mode:
      hashSize = CE2_HASH_SHA384_DIGEST_SIZE_IN_BYTES;
      hashMode = CE2_HASH_SHA384_mode;
      break;
    case CE2_KDF_HASH_SHA512_mode:
      hashSize = CE2_HASH_SHA512_DIGEST_SIZE_IN_BYTES;
      hashMode = CE2_HASH_SHA512_mode;
      break;
    default:
      return CE2_LLF_KDF_MODULE_ERROR_BASE;
  }

  bufferSize = ZZSecretSize + LLF_KDF_COUNTER_SIZE;
  if (OtherInfo_ptr != DX_NULL) {
    if (OtherInfo_ptr->AlgorithmID_ptr != DX_NULL) 
      bufferSize += OtherInfo_ptr->SizeOfAlgorithmID;
    if (OtherInfo_ptr->PartyUInfo_ptr != DX_NULL)
      bufferSize += OtherInfo_ptr->SizeOfPartyUInfo;
    if (OtherInfo_ptr->PartyVInfo_ptr != DX_NULL) 
      bufferSize += OtherInfo_ptr->SizeOfPartyVInfo;
    if (OtherInfo_ptr->SuppPrivInfo_ptr != DX_NULL) 
      bufferSize += OtherInfo_ptr->SizeOfSuppPrivInfo;
    if (OtherInfo_ptr->SuppPubInfo_ptr != DX_NULL) 
      bufferSize += OtherInfo_ptr->SizeOfSuppPubInfo;
  }

  /* d = [keyLen/hashLen] */
  d = (KeyLenInBytes + hashSize -1)/hashSize + 1;

  pHashes = malloc(d*hashSize*sizeof(DxUint8_t));
  pBuffer = malloc(bufferSize*sizeof(DxUint8_t));
  if (pHashes == DX_NULL || pBuffer == DX_NULL) {
    result = CE2_LLF_KDF_MODULE_ERROR_BASE;
    goto error_case;
  }

  /******************************************************/                             
  /* counter = 1                                        */
  /* For i = 1 to d                                     */
  /*   Hash_i = H(ZZ||counter||OtherInfo);              */
  /*   OtherInfo = AlgorithmID, PartyUInfo, PartyVInfo, */
  /*   SuppPrivInfo, SuppPubInfo                        */
  /*   counter = counter + 1                            */ 
  /*   Increment i                                      */
  /******************************************************/                             

  /* Copy ZZSecret to buffer */
  pos = 0;
  memcpy(pBuffer + pos, ZZSecret_ptr, ZZSecretSize);
  pos += ZZSecretSize;
  /* Pass counter position */
  counterPos = pos;
  pos += LLF_KDF_COUNTER_SIZE;
  /* Copy OtherInfo to buffer */
  if (OtherInfo_ptr != DX_NULL) {
    if (OtherInfo_ptr->AlgorithmID_ptr != DX_NULL) {
      memcpy(pBuffer + pos, OtherInfo_ptr->AlgorithmID_ptr, 
        OtherInfo_ptr->SizeOfAlgorithmID);
      pos += OtherInfo_ptr->SizeOfAlgorithmID;
    }
    if (OtherInfo_ptr->PartyUInfo_ptr != DX_NULL) {
      memcpy(pBuffer + pos, OtherInfo_ptr->PartyUInfo_ptr, 
        OtherInfo_ptr->SizeOfPartyUInfo);
      pos += OtherInfo_ptr->SizeOfPartyUInfo;
    }
    if (OtherInfo_ptr->PartyVInfo_ptr != DX_NULL) {
      memcpy(pBuffer + pos, OtherInfo_ptr->PartyVInfo_ptr, 
        OtherInfo_ptr->SizeOfPartyVInfo);
      pos += OtherInfo_ptr->SizeOfPartyVInfo;
    }
    if (OtherInfo_ptr->SuppPubInfo_ptr != DX_NULL) {
      memcpy(pBuffer + pos, OtherInfo_ptr->SuppPubInfo_ptr, 
        OtherInfo_ptr->SizeOfSuppPubInfo);
      pos += OtherInfo_ptr->SizeOfSuppPubInfo;
    }
    if (OtherInfo_ptr->SuppPrivInfo_ptr != DX_NULL) {
      memcpy(pBuffer + pos, OtherInfo_ptr->SuppPrivInfo_ptr, 
        OtherInfo_ptr->SizeOfSuppPrivInfo);
      pos += OtherInfo_ptr->SizeOfSuppPrivInfo;
    }
  }

  i_counter = 1;
  for (i = 0; i < d; i++) {
    /* Insert counter into the buffer */
    /* Convert i_counter to big-endian format */
    pBuffer[counterPos + 0] = (DxUint8_t)(i_counter >> 24);
    pBuffer[counterPos + 1] = (DxUint8_t)(i_counter >> 16);
    pBuffer[counterPos + 2] = (DxUint8_t)(i_counter >> 8);
    pBuffer[counterPos + 3] = (DxUint8_t)(i_counter);

    //PrintBuffer(stdout, pBuffer, bufferSize, " \n\nBefore hash:");

    /* Hash_i = H(ZZ||counter||OtherInfo) */
    error = CE2_HASH(hashMode, pBuffer, bufferSize, hashResult);
    if (error != CE2_OK) {
      result = error;
      goto error_case;
    }

    memcpy(pHashes + i*hashSize, hashResult, hashSize);
    i_counter++;
  }

  /* KeyingData = leftmost keylen bits of Hash1||Hash2||…||Hashd */
  memcpy(KeyingData_ptr, pHashes, KeyLenInBytes);

error_case:
  if (pHashes != DX_NULL)
    free(pHashes);
  if (pBuffer != DX_NULL)
    free(pBuffer);
  return result;
} /* End of LLF_KDF_ConcatKeyDerivFunc */

/************************ Public Functions ********************/

/**
****************************************************************
* Function Name: 
*  LLF_KDF_KeyDerivFunc
*
*  @param[in] ZZSecret_ptr 	- A pointer to shared secret value octet string.
*  @param[in] ZZSecretSize  - The shared secret key Size, in bytes.
*  @param[in] OtherInfo     - The pointer to structure, containing pointers 
*                             and sizes of optional data shared by two
*                             entities intended to share the secret value. 
*                             This argument (structure) and also its members 
*                             are optional (if any is not need - set its pointer 
*                             and size to NULL).
*  @param[in] KDFhashMode	  - The hash function to be used. The hash function 
*                             output must be at least 160 bits.
*  @param[out] KeyingData_ptr - A pointer to the keying data derived from the 
*                               secret key, of length KeyLenInBits
*  @param[in] KeyLenInBytes	 - The size in bytes of the keying data to be generated. 
*                              In our implementation - KeyLenInBytes <= 2^32-1  
*                              (in standards KeyLengthInBits < hashlen*(2^32-1)).
*  @param[in] derivation_mode - Specifies one of above described derivation modes.
*
* @returns \b
*  CE2Error_t 
*  - CE2_OK - on success
*  - Otherwise - on failure a value from CE2_error.h:
*
* \brief \b 
* Description:
*  LLF_KDF_KeyDerivFunc performs key derivation according to one of some modes 
*  defined in standards: ANSI X9.42-2001, ANSI X9.63, OMA_TS_DRM_DRM_V2_0-20050712-C.
*  The present implementation of the function allows the following operation modes:
*  - CE2_KDF_ASN1_DerivMode - mode based on  ASN.1 DER encoding;
*  - CE2_KDF_ConcatDerivMode - mode based on concatenation;
*  - CE2_KDF_X963_DerivMode = CE2_KDF_ConcatDerivMode;
*  - CE2_KDF_OMADRM_DerivMode - specific mode for OMA DRM.
*  The purpose of this function is to derive a keying data from the shared secret 
*  value and some other optional shared information (SharedInfo).
*  The actual APIs that will be used by the user are:
*  - CE2_KDF_ASN1_KeyDerivFunc ;
*  - CE2_KDF_ConcatKeyDerivFunc ;
*  - CE2_KDF_OMADRM_KeyDerivFunc .
*  
*  \note The length in Bytes of the hash result buffer is denoted by "hashlen".
*  \note All buffers arguments are represented in Big-Endian format.
*
*  \b 
* Algorithm:
*  -# Choose derivation_mode and call corresponding subfunction 
***************************************************************/
CE2Error_t  LLF_KDF_KeyDerivFunc(DxUint8_t                *ZZSecret_ptr,
                                 DxUint32_t                ZZSecretSize,
                                 CE2_KDF_OtherInfo_t     *OtherInfo_ptr,
                                 CE2_KDF_HASH_OpMode_t    KDFhashMode,
                                 CE2_KDF_DerivFuncMode_t  derivation_mode,
                                 DxUint8_t                *KeyingData_ptr,
                                 DxUint32_t                KeyLenInBytes )
{
  switch(derivation_mode) {
    case CE2_KDF_ASN1_DerivMode:
      return LLF_KDF_ASN1_KeyDerivFunc(ZZSecret_ptr, ZZSecretSize,
        OtherInfo_ptr, KDFhashMode, KeyingData_ptr, KeyLenInBytes);
/* Disabled fro the time being*/
/*   
    case CE2_KDF_ConcatDerivMode:
      return LLF_KDF_ConcatKeyDerivFunc(ZZSecret_ptr,ZZSecretSize,
        OtherInfo_ptr, KDFhashMode, KeyingData_ptr, KeyLenInBytes);
      break;
*/
    case CE2_KDF_OMADRM_DerivMode:
      return LLF_KDF_ConcatKeyDerivFunc(ZZSecret_ptr, ZZSecretSize,
        DX_NULL, CE2_KDF_HASH_SHA1_mode, KeyingData_ptr, KeyLenInBytes);
      break;
    default:
      return CE2_LLF_KDF_MODULE_ERROR_BASE;
  }

  return CE2_OK;
} /* End of LLF_KDF_KeyDerivFunc */
