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

/************* Include Files ***********************************/
#include "DX_VOS_Mem.h"

#include "CRYS.h"
#include "CRYS_KDF.h"

#include "KMNG_Defs.h" 
#include "KMNG_API.h"

#include "ODRM_TLK_ATP_Defines.h"
#include "ODRM_TLK_ATP_data.h" 
#include "ODRM_TLK_ATP_UTIL.h"

#include "tlk_odrm_api.h"
#include "MW_ATP_UTIL_funcs.h"

/*****************************************************************************
* Function Name:
*  TST_ATP_TLK_ODRM_KmacAndKmekEncryptUnderKri
* 
* Inputs:
*  None
*
* Outputs:
*  DxError_t - Function error return
*
* Description:
*    Test TLK_ODRM_KmacAndKmekEncryptUnderKri
*  
* Algorithm:
* 1. Get work space size
* 2. Initialize a Key Ring  
* 3. insert Kmac and Kmek to the Key Ring 
* 4. Call TLK_ODRM_KmacAndKmekEncryptUnderKri ,returns C
* 5. Extract Kmac and Kmek from C by using the scheme in OMA DRM v2, section 7.2.4
* 6. Validate this is the same Kmac and Kmek we crated C with
*******************************************************************************/

DxError_t TST_ATP_TLK_ODRM_KmacAndKmekEncryptUnderKri(void)
{

  TLK_ODRM_KMNGKeyRing_t      keyesRing;
  KMNG_UserSpecificKeyData_t  UserSpecificKeyData = {0,0};
  TLK_ODRM_KMNGKey_t          kMac;
  DxByte_t     *              kMacBuff = TST_gKekkey;
  TLK_ODRM_KMNGKey_t          kMek;
  DxByte_t   *                kMekBuff = TST_gKeyData;
  DxUint32_t                  kMacAndkMekBuffSize = ODRMTLK_TST_DOUBLE_AES_KEY_SIZE;
  TLK_ODRM_Buffer_t           TLKc = {0x00};
  DxError_t                   TST_Error;
  DxByte_t *                  KEKkey  = TST_gEncryptedDataOut;
 

  TST_gNumOfSymKeys = 2;
  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_KmacAndKmekEncryptUnderKri 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_KmacAndKmekEncryptUnderKri",
                                                 "TST_ATP_TLK_ODRM_KmacAndKmekEncryptUnderKri");

   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_KmacAndKmekEncryptUnderKri",
                                                 "TST_ATP_TLK_ODRM_KmacAndKmekEncryptUnderKri");
 


   /* Build RSA  private key */
   TST_Error = CRYS_RSA_Build_PrivKey(&TST_gRiPrivKey,
	                                    TST_gPrivExponent,
                                      TST_gPrivExponentSize,
                                      TST_gPubExponent,
                                      TST_gPubExponentSize,
                                      TST_gModulus,
                                      TST_gModulusSize);


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

 

   /* Build RSA  public key */
   TST_Error = CRYS_RSA_Build_PubKey(&TST_giPubKey,
                                     TST_gPubExponent,
                                     TST_gPubExponentSize,
                                     TST_gModulus,
                                     TST_gModulusSize);

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

  

   /* 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");
        
      



  /************************************************************************/
  /*2. Init the key ring                                                  */
  /************************************************************************/
  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_KmacAndKmekEncryptUnderKri",
                                                 "TST_ATP_TLK_ODRM_KmacAndKmekEncryptUnderKri");
       

 
  /************************************************************************/
  /* 3.insert Kmac and Kmek to the Key Ring                               */
  /************************************************************************/
/*Kmek*/
 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, 
                                    kMekBuff,
                                    &kMek.keyId,
                                    TST_gKeyRingBuf);

  
   TST_ATP_CHECK_ERROR_RETURN(TST_Error,"KMNG_ImportSymUserKey kMek",
                                                 "TST_ATP_TLK_ODRM_KmacAndKmekEncryptUnderKri",
                                                 "TST_ATP_TLK_ODRM_KmacAndKmekEncryptUnderKri");
        

 /*Kmac*/
 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, 
                                    kMacBuff,
                                    &kMac.keyId,
                                    TST_gKeyRingBuf);

  
   TST_ATP_CHECK_ERROR_RETURN(TST_Error,"KMNG_ImportSymUserKey kMac",
                                                 "TST_ATP_TLK_ODRM_KmacAndKmekEncryptUnderKri",
                                                 "TST_ATP_TLK_ODRM_KmacAndKmekEncryptUnderKri");
        
   


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

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


  /************************************************************************/
  /* load TLKc parameters                                                 */
  /************************************************************************/
  TLKc.buff_ptr = TST_gC;
  TLKc.buffSizeInBytes = ODRMTLK_TST_RSA_1024_KEY_SIZE_IN_BYTES+ ODRMTLK_TST_C2_SIZE;
    
  /**************************************************************************************************/
  /*4. Call TLK_ODRM_KmacAndKmekEncryptUnderKri                                                       */
  /* encrypts Kmac and Kmek under rights issuer's public key and returns to user encrypted C string */
  /**************************************************************************************************/
   TST_Error = TLK_ODRM_KmacAndKmekEncryptUnderKri(&keyesRing,
                                                   &TST_giPubKey,
                                                   &kMac,
                                                   &kMek,
                                                   TLKc,
                                                   TST_gWorkspace,
                                                   ODRMTLK_TST_WORKSPACE_SIZE);

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

 /*********************************************************************************************/
 /*5.Extract Kmac and Kmek from C by using the scheme in OMA DRM v2, section 7.2.4              */
 /*********************************************************************************************/
 /* C1|C2 = C*/

 /*c1 = OS2IP(C1,mLen)*/
   DX_VOS_FastMemCpy(TST_gC1,TLKc.buff_ptr,ODRMTLK_TST_RSA_1024_KEY_SIZE_IN_BYTES);

 /*Z = RSA.DECRYPT(PrivKeyRI, c1) = c1 d mod m*/
  TST_Error = CRYS_RSA_PRIM_Decrypt(&TST_gRiPrivKey,
                                    &TST_gPrimeData,
                                    TST_gC1,
                                    ODRMTLK_TST_RSA_1024_KEY_SIZE_IN_BYTES,
                                    TST_gZBuff);

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


 /*KEK = KDF(I2OSP(Z, mLen), NULL, kekLen)*/
  TST_Error = CRYS_KDF_OMADRM_KeyDerivFunc(TST_gZBuff,
                                           ODRMTLK_TST_RSA_1024_KEY_SIZE_IN_BYTES,
                                           KEKkey,
                                           ODRMTLK_TST_AES_KEY_SIZE);

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

 /*KMAC | KMEK = AES-UNWRAP(KEK, C2)*/
   DX_VOS_FastMemCpy(TST_gC2Buff, TLKc.buff_ptr+ODRMTLK_TST_RSA_1024_KEY_SIZE_IN_BYTES, ODRMTLK_TST_C2_SIZE );
   TST_Error = CRYS_AES_Unwrap(TST_gC2Buff, 
                               ODRMTLK_TST_C2_SIZE ,
                               KEKkey,
                               CRYS_AES_Key128BitSize,
                               0,
                               TST_gKMacAndkMekBuff,
                               &kMacAndkMekBuffSize);

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

   /*************************************************************/
   /*6. Validate this is the same Kmac and Kmek we crated C with  */
   /************************************************************/
 
  if(0 == DX_VOS_MemCmp(kMacBuff,TST_gKMacAndkMekBuff,ODRMTLK_TST_AES_KEY_SIZE))
     TST_Error = TST_PASS; 
  else
     TST_Error = TST_FAIL; 
  
   TST_ATP_CHECK_ERROR_RETURN(TST_Error,"DX_VOS_MemCmp - kMacBuff",
                                                 "TST_ATP_TLK_ODRM_KmacAndKmekEncryptUnderKri",
                                                 "TST_ATP_TLK_ODRM_KmacAndKmekEncryptUnderKri");
   


  if(0 == DX_VOS_MemCmp(kMekBuff,TST_gKMacAndkMekBuff + ODRMTLK_TST_AES_KEY_SIZE,ODRMTLK_TST_AES_KEY_SIZE))
    TST_Error = TST_PASS; 
  else
    TST_Error = TST_FAIL; 
  
   TST_ATP_CHECK_ERROR_RETURN(TST_Error,"DX_VOS_MemCmp - kMekBuff",
                                                 "TST_ATP_TLK_ODRM_KmacAndKmekEncryptUnderKri",
                                                 "TST_ATP_TLK_ODRM_KmacAndKmekEncryptUnderKri");
   


  ATP_LOG_TST_PRINT((MW_ATP_MSG," **************************************************************\n\n"));
  ATP_LOG_TST_PRINT((MW_ATP_MSG," ** TST_ATP_TLK_ODRM_KmacAndKmekEncryptUnderKri  Test  PASS !!! \n"));     
  ATP_LOG_TST_PRINT((MW_ATP_MSG," ************************************************************** \n \n"));
EXIT_ON_ERROR:
 return TST_Error;

}
