/*! \file 
**********************************************************************************
*Title:                         Discretix OMA DRM v2 Toolkit ATP Test source file
*
* Filename:                     odrm_tlk_ATP_Kmac_And_Kmek_Derive_Under_DomainKey.c
*
* 
* Created:                      05.12.2007
*
*
* \Author                       SagitB
*
* \Remarks
*           Copyright (C) 2007 by Discretix Technologies Ltd. All Rights reserved.
**********************************************************************************/

/************* Include Files ***********************************/
#include "CRYS.h"
#include "tlk_odrm_api.h"
#include "ODRM_TLK_ATP_Defines.h"
#include "KMNG_Defs.h" 
#include "KMNG_API.h"
#include "ODRM_TLK_ATP_data.h" 
#include "CRYS_KDF.h"
#include "CRYS_RSA_KG.h"
#include "DX_VOS_Mem.h"
#include "CRYS_RSA_PRIM.h"
#include "ODRM_TLK_ATP_UTIL.h" 
#include "CRYS_KMNG.h"
#include "MW_ATP_UTIL_funcs.h"
 

/*****************************************************************************
* Function Name:                                                           
*  TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey
* 
* Inputs:
*  None
*
* Outputs:
*  DxError_t - Function error return                                        
*
* Description:                                                                 
*    Test TLK_ODRM_KmacAndKrekDeriveUnderDomainKey
*  
* Algorithm:
* 1.  Get work space size 
*.2   Init the key ring 
* 3.  Insert KD  (domain Key) to the Key ring 
* 4.  Perform  KDF (domain Generation on) KD 
* 5.    Create C with Kd ,Kmac and Krek by using the scheme in OMA DRM v2, section 7.2.3
* 6.    Call TLK_ODRM_KmacAndKrekDeriveUnderDomainKey
* 7.    Activate Kmac and Krek from the key ring using the key ids received from the previous operation
* 8.    Validate this is the same Kmac and Krek we crated C with
*  8.1 preform HMAC with Kmac on plain text
*  8.2 compare the result with plain text that was Hmac with Kmac in advance
*  8.3 for verifying KD we will wrap a Cek key with Kd in advance
*       perform  TLK_ODRM_AesUnwrapAndStore with kd as wrapping Key
*       Encrypt plain text  with Cek key and compare with text we encrypted in advance with Cek
*  8.4 load parameters and callTLK_ODRM_AesUnwrapAndStore
*  8.5 Activate  Cek from the key ring using  the key id received from the previous operation
*  8.6 Validate this is the same key we inserted in TLK_ODRM_AesUnwrapAndStore        
*      by Encrypting   plain text with cek from key ring and comparing it with              
*      Plain text we Encrypt in advance with original Cek

*******************************************************************************/   

DxError_t TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey(void)
{
  TLK_ODRM_KMNGKey_t          KD;
  DxByte_t *                  KDkey = TST_gKey;
  DxUint8_t                   nextKD[ODRMTLK_TST_AES_KEY_SIZE];
  DxUint32_t                  nextKDSize = ODRMTLK_TST_AES_KEY_SIZE;
  TLK_ODRM_KMNGKey_t          kMac;
  KMNG_HMAC_WrappedKey_t      kMac_WrappedKey;
  TLK_ODRM_KMNGKey_t          kRek;
  KMNG_AES_WrappedKey_t       kRek_WrappedKey;
  TLK_ODRM_KMNGKey_t          TLKCekKey;
  KMNG_AES_WrappedKey_t       Cek_WrappedKey;
  TLK_ODRM_Buffer_t           TLKwrappedCekKeyBuff;
  DxError_t                   TST_Error;
  DxUint32_t                  cSize = ODRMTLK_TST_CBUFF_SIZE;
  TLK_ODRM_Buffer_t           c;
  TLK_ODRM_KMNGKeyRing_t      keyesRing;
  KMNG_UserSpecificKeyData_t  UserSpecificKeyData = {0,0};
  DxUint32_t                  index;
  TLK_ODRM_DomainGeneration_t domainGeneration;
  DxByte_t*                   EncryptedPlainBuffWithCek   = TST_gEncryptedPlainBuffWithK;
  
  TST_gNumOfSymKeys = 4;
  TST_gNumOfRSAKeys = 0;
  TST_gNumOfDHKeys  = 0;
   
  ATP_LOG_TST_PRINT((MW_ATP_MSG," ********************************************************* \n"));
  ATP_LOG_TST_PRINT((MW_ATP_MSG," ** TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey Test \n"));     
  ATP_LOG_TST_PRINT((MW_ATP_MSG," ******************************************************** \n \n"));

  
  #ifdef  ODRM_TLK_CLEAR_KEY_RING
    DX_VOS_MemSetZero(TST_gKeyRingBuf,ODRMTLK_TST_MAX_SIZE_KEY_RING_BUF);
    DX_VOS_MemSetZero(TST_gDevKeyRingBuf,ODRMTLK_TST_MAX_SIZE_DEV_KEY_RING_BUF);
  #endif


  /* 1.Get work space size */
  TST_Error = TLK_ODRM_WorkspaceSizeGet(TLK_ODRM_MIN_WORKSPACE,&TST_gWorkspaceSizeInBytes);

  TST_ATP_CHECK_ERROR_RETURN(TST_Error,"TLK_ODRM_WorkspaceSizeGet",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey");
 
  if(TST_gWorkspaceSizeInBytes < ODRMTLK_TST_WORKSPACE_SIZE)
    TST_Error = TST_FAIL;
  else
    TST_Error = TST_PASS;
  
   
  TST_ATP_CHECK_ERROR_RETURN(TST_Error,"Insufficient Workspace Size ",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey");
 

  /************************************************************************/
  /* 2.Init the key ring                                                  */
  /************************************************************************/

  /* Get key ring expected buffer size */
  TST_Error = KMNG_GetKeyRingBufferSize(TST_gNumOfSymKeys, 
                                        TST_gNumOfRSAKeys, 
                                        TST_gNumOfDHKeys, 
                                        &TST_gRequiredBufSizeInBytes); 

  TST_ATP_CHECK_ERROR_RETURN(TST_Error,"KMNG_GetKeyRingBufferSize",
                                                "TST_ATP_TLK_ODRM_KmacAndKmekEncryptUnderKri",
                                                "TST_ATP_TLK_ODRM_KmacAndKmekEncryptUnderKri");

 
  TST_Error = KMNG_KeyRingInit(TST_gPasswordSize,
                               TST_gPassword,
                               TST_gNumOfSymKeys, 
                               TST_gNumOfRSAKeys,
                               TST_gNumOfDHKeys, 
                               TST_gKeyRingBuf);

  TST_ATP_CHECK_ERROR_RETURN(TST_Error,"KMNG_KeyRingInit",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey");

  


  /************************************************************************/
  /*3. Insert KD  (domain Key) to the Key ring                            */
  /************************************************************************/

  /* User Specific KeyData - insert the number of domain generation */
  UserSpecificKeyData[0] = ODRMTLK_TST_MAX_GENERATION_NUMBER;

  /*KD*/
  TST_Error = KMNG_ImportSymUserKey( TST_gPassword,
                                     TST_gPasswordSize,
                                     TST_gPassword,
                                     TST_gPasswordSize,
                                     KMNG_KeyTypeAES,
                                     ODRMTLK_TST_AES_KEY_SIZE,
                                     KMNG_KEY_USAGE_ODRM,
                                     TLK_ODRM_KEY_RESTRICTION,
                                     UserSpecificKeyData, 
                                     KDkey,
                                     &KD.keyId ,
                                     TST_gKeyRingBuf);

  TST_ATP_CHECK_ERROR_RETURN(TST_Error,"KMNG_ImportSymUserKey KD",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey");



  /*******************************************/
  /* 4.perform  KDF (domain Generation on) KD */
  /******************************************/
  domainGeneration = ODRMTLK_TST_CURRENT_GENERATION_NUMBER;

  for (index = 0; index < ODRMTLK_TST_MAX_GENERATION_NUMBER - ODRMTLK_TST_CURRENT_GENERATION_NUMBER ; index++)
  {
    TST_Error = CRYS_KDF_OMADRM_KeyDerivFunc(KDkey,
                                             ODRMTLK_TST_AES_KEY_SIZE,
                                             nextKD,
                                             nextKDSize);

    TST_ATP_CHECK_ERROR_RETURN(TST_Error,"CRYS_KDF_OMADRM_KeyDerivFunc",
                                                  "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey",
                                                  "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey");

    

    DX_VOS_FastMemCpy(KDkey, nextKD, nextKDSize);   

  }

  /*********************************************************************************************/
  /*5. Create C with Kd ,Kmac and Krek by using the scheme in OMA DRM v2, section 7.2.3        */
  /*********************************************************************************************/

  /*C = AES-WRAP(KD, KMAC | KREK)*/
  TST_Error = CRYS_AES_Wrap(TST_gKMacAndkRekBuff,
                            ODRMTLK_TST_DOUBLE_AES_KEY_SIZE,
                            KDkey,
                            CRYS_AES_Key128BitSize,
                            0,
                            TST_gCBuff,
                            &cSize );

  TST_ATP_CHECK_ERROR_RETURN(TST_Error,"CRYS_AES_Wrap",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey");


  /************************************************************************/
  /* load key ring parameters                                             */
  /************************************************************************/
  keyesRing.keyKmngRing_ptr               = TST_gKeyRingBuf;
  keyesRing.keyKmngRingPassword_ptr       = TST_gPassword;
  keyesRing.keyKmngRingPasswordLenInBytes = TST_gPasswordSize;


  /************************************************************************/
  /* load KD parameters                                                   */
  /************************************************************************/
  KD.keyPassword_ptr       = TST_gPassword;
  KD.keyPasswordLenInBytes = TST_gPasswordSize;

  /************************************************************************/
  /* load c parameters                                                    */
  /************************************************************************/
  c.buff_ptr = TST_gCBuff;
  c.buffSizeInBytes = cSize;

  /************************************************************************/
  /* load kMac parameters                                                 */
  /************************************************************************/
  kMac.keyPassword_ptr = TST_gPassword;
  kMac.keyPasswordLenInBytes = TST_gPasswordSize;

  /************************************************************************/
  /* load kRek parameters                                                 */
  /************************************************************************/
  kRek.keyPassword_ptr = TST_gPassword;
  kRek.keyPasswordLenInBytes = TST_gPasswordSize;

  /**************************************************************************************************/
  /* 6.Call TLK_ODRM_KmacAndKrekDeriveUnderDomainKey                                                */
  /* gets from user an encrypted c-string used for Kmac and Krek derivation and domain generation   */
  /* unwraps them using given domain key and stores  the keys in the key ring provided by the user  */
  /* returns id for Kmac and KRek                                                                   */
  /**************************************************************************************************/
  TST_Error = TLK_ODRM_KmacAndKrekDeriveUnderDomainKey(&keyesRing,
                                                       &KD,
                                                       domainGeneration,
                                                       c,
                                                       &kMac,
                                                       &kRek,
                                                       TST_gWorkspace,
                                                       ODRMTLK_TST_WORKSPACE_SIZE);

  TST_ATP_CHECK_ERROR_RETURN(TST_Error,"TLK_ODRM_KmacAndKrekDeriveUnderDomainKey",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey");

  /*****************************************************************************************************/
  /* 7. Activate Kmac and Krek from the key ring using the key ids received from the previous operation  */
  /*****************************************************************************************************/

  /*prepares the kMac key for CRYS function */

  TST_Error =  KMNG_ActivateHMACKey(keyesRing.keyKmngRingPassword_ptr,
                                    keyesRing.keyKmngRingPasswordLenInBytes,
                                    kMac.keyPassword_ptr, 
                                    kMac.keyPasswordLenInBytes,
                                    kMac.keyId,
                                    keyesRing.keyKmngRing_ptr,
                                    kMac_WrappedKey );


  TST_ATP_CHECK_ERROR_RETURN(TST_Error,"KMNG_ActivateHMACKey  - kMac",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey");


  /*prepares the kRek key for CRYS function */
  TST_Error = KMNG_ActivateAESKey(keyesRing.keyKmngRingPassword_ptr,
                                  keyesRing.keyKmngRingPasswordLenInBytes,
                                  kRek.keyPassword_ptr, 
                                  kRek.keyPasswordLenInBytes,
                                  kRek.keyId,
                                  keyesRing.keyKmngRing_ptr,
                                  kRek_WrappedKey );


  TST_ATP_CHECK_ERROR_RETURN(TST_Error,"KMNG_ActivateAESKey  - kMac",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey");

  

  /*****************************************************************/
  /* 8.  Validate this is the same Kmac and Krek we crated C with    */
  /*****************************************************************/
  /*kMac*/
  /*8.1 preform HMAC with Kmac on plain text*/
  TST_Error = CRYS_KMNG_HMAC(kMac_WrappedKey,     
                             CRYS_HASH_SHA1_mode,
                             TST_gPlainBuff,
                             ODRMTLK_TST_AES_KEY_SIZE,
                             TST_gHmacResultBuff);

  TST_ATP_CHECK_ERROR_RETURN(TST_Error,"CRYS_KMNG_HMAC  - kMac",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey");

 

  /* 8.2 compare the result with plain text that was Hmac with Kmac in advance*/
  if(0 != DX_VOS_MemCmp(TST_gHmacPlainBuffWithKmac,TST_gHmacResultBuff,ODRMTLK_TST_AES_KEY_SIZE))
    TST_Error = TST_FAIL; 
  else
    TST_Error = TST_PASS; 

  TST_ATP_CHECK_ERROR_RETURN(TST_Error,"DX_VOS_MemCmp - kMac",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey");

  


  /* 8.3 for verifying Krek we will wrap a Cek key with Kd in advance*/
  /*perform  TLK_ODRM_AesUnwrapAndStore with Krek as wrapping Key*/
  /*Encrypt plain text  with Cek key and compare with text we encrypted in advance with Cek*/
  /*Krek */

  /*****************************************************/
  /*8.4 load parameters and callTLK_ODRM_AesUnwrapAndStore */
  /******************************************************/ 


  /*Krek */

  /************************************************************************/
  /* load wrapped Cek  buff                                                 */
  /************************************************************************/
  TLKwrappedCekKeyBuff.buff_ptr              = TST_gWrappedCekWithRek;
  TLKwrappedCekKeyBuff.buffSizeInBytes       = ODRMTLK_TST_WRAPPED_KEY_LENGTH_IN_BYTES;

  /************************************************************************/
  /* load  Cek  pass parameters                                             */
  /************************************************************************/
  TLKCekKey.keyPassword_ptr               = TST_gPassword;
  TLKCekKey.keyPasswordLenInBytes         = TST_gPasswordSize;


  TST_Error = TLK_ODRM_AesUnwrapAndStore(&keyesRing,
                                         &kRek,
                                         &TLKwrappedCekKeyBuff,
                                         &TLKCekKey,
                                         TST_gWorkspace,
                                         ODRMTLK_TST_WORKSPACE_SIZE);


  TST_ATP_CHECK_ERROR_RETURN(TST_Error,"TLK_ODRM_AesUnwrapAndStore  - kRek",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey");

  

  /*****************************************************************************************************/
   /* 8.5 Activate  Cek from the key ring using  the key id received from the previous operation    */
   /*************************************************************************************************/    


  /*prepares the Cek key for CRYS function */
  TST_Error = KMNG_ActivateAESKey(TST_gPassword,
                                  TST_gPasswordSize,
                                  TLKCekKey.keyPassword_ptr, 
                                  TLKCekKey.keyPasswordLenInBytes,
                                  TLKCekKey.keyId,
                                  TST_gKeyRingBuf,
                                  Cek_WrappedKey );


  TST_ATP_CHECK_ERROR_RETURN(TST_Error,"KMNG_ActivateAESKey Cek",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey");



  /* 8.6 Validate this is the same key we inserted in TLK_ODRM_AesUnwrapAndStore
  by Encrypting   plain text with cek from key ring and comparing it with
  Plain text we Encrypt in advance with original Cek*/

  TST_Error = CRYS_KMNG_AES(Cek_WrappedKey,
                            TST_gIvCounter,       
                            CRYS_AES_Encrypt, 
                            CRYS_AES_ECB_mode,       
                            TST_gPlainBuff,        
                            ODRMTLK_TST_AES_KEY_SIZE,         
                            TST_gEncryptedDataOut); 

  TST_ATP_CHECK_ERROR_RETURN(TST_Error,"CRYS_KMNG_AES  - Cek Key",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey");

 

  if(0 != DX_VOS_MemCmp(EncryptedPlainBuffWithCek,TST_gEncryptedDataOut,ODRMTLK_TST_AES_KEY_SIZE))
    TST_Error =TST_FAIL; 
  else
   TST_Error = TST_PASS; 

  TST_ATP_CHECK_ERROR_RETURN(TST_Error,"DX_VOS_MemCmp - Cekkey",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey",
                                                "TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey");

  

  ATP_LOG_TST_PRINT((MW_ATP_MSG," ***************************************************************** \n\n"));
  ATP_LOG_TST_PRINT((MW_ATP_MSG," ** TST_ATP_TLK_ODRM_KmacAndKrekDeriveUnderDomainKey  Test  PASS !!! \n"));     
  ATP_LOG_TST_PRINT((MW_ATP_MSG," **************************************************************** \n \n"));

EXIT_ON_ERROR:
  return TST_Error;
}
