/**@@@+++@@@@******************************************************************
**
** Microsoft Windows Media
** Copyright (C) Microsoft Corporation. All rights reserved.
**
***@@@---@@@@******************************************************************
*/

//#include <stdlib.h>
//#include <stdio.h>
// #include <string.h>
// #include <math.h>
#include "DX_VOS_String.h"
#include "DX_VOS_Utils.h"
#include "tclib.h"

#include "drmpkcrypto.h"
#include "drmcipher.h"
#include "oemimpl.h"
#include "drmutilities.h"
#include "drmsha1.h"
#include "drmhmac.h"
#include "tOEMIMP.h"

//#define DRM_PKCRYPTO_CONTEXT_BUFFER_SIZE_LESS 1
//#define DRM_PKCRYPTO_CONTEXT_BUFFER_SIZE_MORE 399360
//DRM_BYTE rgbCryptoContextLessSize[DRM_PKCRYPTO_CONTEXT_BUFFER_SIZE_LESS];   /* size of the buffer pkcrypto code needs. Note this may have alignment problems. */
//DRM_BYTE rgbCryptoContextMoreSize[DRM_PKCRYPTO_CONTEXT_BUFFER_SIZE_MORE];   /* size of the buffer pkcrypto code needs. Note this may have alignment problems. */

DRM_BYTE g_szPlainKey[PK_ENC_PLAINTEXT_LEN] = {"PlainKy"};/*{7, "PlainKy"}*/; /*max of 7 - export limit */
DRM_BYTE* g_szKey;
DRM_CRYPTO_CONTEXT oCrypto;
DRM_BYTE* prgbCryptoContext;   /* size of the buffer pkcrypto code needs. Note this may have alignment problems. */
DRM_BYTE        cipher    [PK_ENC_CIPHERTEXT_LEN]; 
char tmp[PK_ENC_CIPHERTEXT_LEN+1];

DRM_BYTE bRandBuf[PK_ENC_PLAINTEXT_LEN];
DRM_BYTE bOutBuf[PK_ENC_CIPHERTEXT_LEN];
DRM_BYTE bDecryptBuf[PK_ENC_PLAINTEXT_LEN];
DRM_BYTE *bRandBufIn;
DRM_BYTE *bCipher;
DRM_BYTE *bDecryptBufOut;
PUBKEY *pubKey;
PRIVKEY *privKey;

const DRM_CHAR g_szClearText[] = "Sample clear text.  DRM content encryption will be applied.";
DRM_BYTE g_rgbCipherText[sizeof(g_szClearText)] =
{
 0x77, 0xd7, 0x89, 0xf2, 0x1a, 0xf5, 0x24, 0x30, 0x2d, 0xe6, 0x9c, 0x28, 0x04, 
 0x36, 0x70, 0x28, 0xd9, 0xa4, 0xd0, 0xe3, 0x6e, 0x46, 0x65, 0x3f, 0x03, 0xee, 
 0xf6, 0x95, 0xfb, 0xe3, 0x58, 0xef, 0xed, 0xa5, 0x01, 0xfa, 0x21, 0x32, 0x06,
 0x6b, 0xa3, 0x73, 0x86, 0xc8, 0xb2, 0x78, 0xc2, 0x20, 0xab, 0x6a, 0xf1, 0x63,
 0x64, 0x81, 0x13, 0x4d, 0x28, 0xc2, 0x7d, 0x1c
};

DRM_BYTE g_rgbContentKey[DRMCIPHERKEYLEN] = { 1,2,3,4,5,6,7 };
DRM_CIPHER_CONTEXT oContext;
DRM_CIPHER_CONTEXT *pContext;
DRM_BYTE *rgbEncData;
DRM_BYTE *rgbClearData;
SHA_CONTEXT g_oShaContext;
HMAC_CONTEXT gHmacContext = {0};
RC4_KEYSTRUCT rc4ks = {0};
DESTable  desTable = {0};
/* CBC Hashing */
DRM_BYTE *pbOrgData;
DRM_BYTE *pbEncData;
DRM_CBCKey m_mackey, m_invmackey; 

#define MAX_DATA_COUNT 6
DRM_BYTE *g_pbRandomData[MAX_DATA_COUNT];

DRM_RESULT TestGenRandomData(long argc, char **argv)
{
    DRM_RESULT dr = DRM_SUCCESS;
    DRM_DWORD i, cbData;

    ChkArg(argc == 2 && argv[0] && argv[1]);

	   i = OEM_atol(argv[0]);
    ChkArg(i < MAX_DATA_COUNT);

    cbData = OEM_atol(argv[1]);

    DX_VOS_MemFree(g_pbRandomData[i]);
    ChkMem(g_pbRandomData[i] = (DRM_BYTE*)DX_VOS_MemMalloc(cbData));
    OEM_GenRandomBytes(g_pbRandomData[i], cbData);
ErrorExit:
    return dr;    
}

DRM_RESULT TestMemComp(long argc, char **argv)
{
    DRM_RESULT dr = DRM_SUCCESS;
    DRM_DWORD i, j;

    ChkArg(argc == 3 && argv[0] && argv[1] && argv[2]);

    i = OEM_atol(argv[0]);
    ChkArg(i < MAX_DATA_COUNT);
	j = OEM_atol(argv[1]);
    ChkArg(j < MAX_DATA_COUNT);

    if (DX_VOS_MemCmp(g_pbRandomData[i], g_pbRandomData[j], OEM_atol(argv[2])))
        dr = DRM_S_FALSE;
ErrorExit:
    return dr;
}

DRM_RESULT TestPKEncrypt_PKDecrypt(long argc, char **argv)
{
    DRM_RESULT dr = DRM_SUCCESS;
    PUBKEY *pubKey2=NULL;
    PRIVKEY *privKey2=NULL;
    PUBKEY *tmpPubKey=NULL;
    PRIVKEY *tmpPrivKey=NULL;

    
    if(argv[0]&&!DX_VOS_StrCmp(argv[0],"GENKEYPAIR"))
    {    
        if(!argv[1])
            prgbCryptoContext = NULL;
        if(!argv[2])
            pubKey = NULL;
        if(!argv[3])
            privKey = NULL;
    }

    ChkDR( DRM_PK_GenKeyPair( prgbCryptoContext,
                        pubKey, 
                        privKey) );

    tmpPubKey = pubKey;
    tmpPrivKey = privKey;

    pubKey2 = (PUBKEY*)OEM_malloc(PK_ENC_PUBLIC_KEY_LEN  );
    privKey2 = (PRIVKEY*)OEM_malloc(PK_ENC_PRIVATE_KEY_LEN  );
    ChkDR( DRM_PK_GenKeyPair( oCrypto.rgbCryptoContext,
                                pubKey2, 
                                privKey2) );

    OEM_GenRandomBytes(bRandBuf, sizeof(bRandBuf ));

    if(argv[0]&&!DX_VOS_StrCmp(argv[0],"ENCRYPT"))
    {    if(!argv[1])
            prgbCryptoContext = NULL;
        
        if(argv[2]&&!DX_VOS_StrCmp(argv[2],"PubKey2"))
        {
            tmpPubKey = pubKey2;
        }
        else if(!argv[2])
            tmpPubKey = NULL;
        if(!argv[3])
            bRandBufIn = NULL;
        if(!argv[4])
            bCipher = NULL;
    }
    dr = DRM_PK_Encrypt( prgbCryptoContext,
                    tmpPubKey, 
                    bRandBufIn, 
                    bCipher );

    if ( dr != DRM_SUCCESS )
    {
        Log("ERROR", "test1.cpp: PKencrypt failed");
        goto ErrorExit;        
    }

    if(argv[0]&&!DX_VOS_StrCmp(argv[0],"DECRYPT"))
    {    if(!argv[1])
            prgbCryptoContext = NULL;
        
        if(argv[2]&&!DX_VOS_StrCmp(argv[2],"PrivKey2"))
            tmpPrivKey = privKey2;
        else if (!argv[2])
            tmpPrivKey = NULL;
        if(!argv[3])
            bCipher = NULL;
        if(!argv[4])
            bDecryptBufOut = NULL;
    }
    dr = DRM_PK_Decrypt(prgbCryptoContext, 
                    tmpPrivKey, 
                    bCipher, 
                    bDecryptBufOut );

    if ( dr != DRM_SUCCESS )
    {
        Log("ERROR", "test1.cpp: PKdecrypt failed");
        goto ErrorExit;        
    }

    if ( DX_VOS_MemCmp( bRandBufIn, bDecryptBufOut, PK_ENC_PLAINTEXT_LEN ) != 0 ) 
    {
        Log("ERROR", "PK encr/decr failed."); 
        dr = DRM_E_FAIL;
        goto ErrorExit;
    }
    DX_VOS_FastMemCpy(tmp, bDecryptBufOut, PK_ENC_PLAINTEXT_LEN);
    Log("Trace", "Decrypted Text = %s", tmp);

ErrorExit:
    OEM_free(pubKey2);
    OEM_free(privKey2);
    return dr;
}

DRM_RESULT TestPKSign_PKVerify(long argc, char **argv)
{
    DRM_RESULT dr = DRM_SUCCESS;
    DRM_DWORD keyLength = 0;
    DRM_BYTE* bSignature;
    DRM_BYTE bSign[PK_ENC_SIGNATURE_LEN];
    PUBKEY *pubKey2 = NULL;
    PRIVKEY *privKey2 = NULL;
    PUBKEY *tmpPubKey;
    PRIVKEY *tmpPrivKey;

    bSignature = bSign;
    ChkDR( DRM_PK_GenKeyPair( prgbCryptoContext,
                        pubKey, 
                        privKey) );

    tmpPubKey = pubKey;
    tmpPrivKey = privKey;

    pubKey2 = (PUBKEY*)OEM_malloc(PK_ENC_PUBLIC_KEY_LEN  );
    privKey2 = (PRIVKEY*)OEM_malloc(PK_ENC_PRIVATE_KEY_LEN  );
    ChkDR( DRM_PK_GenKeyPair( oCrypto.rgbCryptoContext,
                                pubKey2, 
                                privKey2) );

    if(argv[0]&&!DX_VOS_StrCmp(argv[0],"PKSIGN"))
    {    
        if(!argv[1])
            prgbCryptoContext = NULL;        
        if(argv[2]&&!DX_VOS_StrCmp(argv[2],"PrivKey2"))
        {
            tmpPrivKey = privKey2;
        }
        else if (!argv[2])
            tmpPrivKey = NULL;
        if(!argv[3])
            g_szKey = NULL;
        if(argv[3]&&!DX_VOS_StrCmp(argv[4],"PlainKeyLength"))
            keyLength = (DRM_DWORD)DX_VOS_StrLen( (const DxChar*)g_szKey );
        else    /*if(argv[4]&&!strcmp(argv[4],"PlainKeyZeroLength"))*/
            keyLength = 0;
        if(!argv[5])
            bSignature = NULL;
    }
    else
        keyLength = (DRM_DWORD)DX_VOS_StrLen( (const DxChar*)g_szKey );

    ChkDR( DRM_PK_Sign( prgbCryptoContext, tmpPrivKey, g_szKey, keyLength, bSignature) );

    if(argv[0]&&!DX_VOS_StrCmp(argv[0],"PKVERIFY"))
    {    
        if(!argv[1])
            prgbCryptoContext = NULL;        
        if(argv[2]&&!DX_VOS_StrCmp(argv[2],"PubKey2"))
        {
            tmpPubKey = pubKey2;
        }
        else if(!argv[2])
            tmpPubKey = NULL;
        if(!argv[3])
            g_szKey = NULL;
        if(argv[3]&&!DX_VOS_StrCmp(argv[4],"PlainKeyLength"))
            keyLength = (DRM_DWORD)DX_VOS_StrLen( (const DxChar*)g_szKey );
        else    /*if(argv[4]&&!strcmp(argv[4],"PlainKeyZeroLength"))*/
            keyLength = 0;
        if(!argv[5])
            bSignature = NULL;
    }

    if(!DRM_PK_Verify(prgbCryptoContext, tmpPubKey, g_szKey, keyLength, bSignature))
    {    
        Log("Trace", "error in PKsign/verify detected");

        dr = DRM_E_FAIL;
        goto ErrorExit;
    }

ErrorExit:
    OEM_free(pubKey2);
    OEM_free(privKey2);
    return dr;
}

/*
*   TestDRMCipher_Init
*    argv[0]:Length in bytes of the content key
*    argv[1]:Array of bytes of the content key 
*    argv[2]:Pointer to a DRM_CIPHER_CONTEXT structure 
*/

DRM_RESULT TestDRMCipher_Init(long argc, char **argv)
{
    DRM_RESULT dr = DRM_SUCCESS;
    DRM_DWORD cbContentKey;
    DRM_BYTE* pbContKey = NULL;

    ChkArg(argc == 3 && argv[0]);

    /* set content key length */
    cbContentKey = OEM_atoi(argv[0]);

    /* set up content key parameter */
    if (argv[1]) {
        ChkMem(pbContKey = (DRM_BYTE*)OEM_malloc(cbContentKey+ 1));
        DX_VOS_MemSet(pbContKey, 0, cbContentKey );
        if (DX_VOS_FindStr( argv[1], "BlobContentKey" ))
            OEM_GenRandomBytes(pbContKey, cbContentKey);
    }


    /* set up cipher context */
    pContext = argv[2]? &oContext: NULL;

    ChkDR(DRM_CPHR_Init(pContext, cbContentKey, pbContKey));

ErrorExit:
    OEM_free(pbContKey);
    return dr;
}

/*
*    TestDRMCipher_Encrypt
*    argv[0]:status of crypto context: NULL or NORMAL
*    argv[1]:Number of bytes to be encrypted by the cipher
*    argv[2]:optional status of pbData: NULL or NORMAL
*/
DRM_RESULT TestDRMCipher_Encrypt(long argc, char **argv)
{
    DRM_RESULT dr = DRM_SUCCESS;
    DRM_DWORD cbData = 0;
    DRM_BYTE* pbData = NULL;

    ChkArg(argc > 1 && argv[1]);

    /* set content key length */
    cbData = OEM_atol(argv[1]);

    /* set up content key to be encrypted */
    if (argc < 3 || argv[2]) {
        ChkMem(pbData = (DRM_BYTE*)OEM_malloc(cbData));
        OEM_GenRandomBytes(pbData, cbData );
        DX_VOS_FastMemCpy( rgbClearData, pbData, cbData);
    }

    ChkDR(DRM_CPHR_Encrypt(argv[0]? &oContext: NULL, cbData, pbData));
    if (pbData) {
        if (!DX_VOS_MemCmp( pbData, rgbClearData, cbData ))
        {
            Log("ERROR", "CryptoTest.cpp: DRMCipher_Encrypt failed");
            ChkDR(DRM_E_FAIL);
        }
        DX_VOS_FastMemCpy( rgbEncData, pbData, cbData);
    }

ErrorExit:
    OEM_free(pbData);
    return dr;
}

/*
* TestDRMCipher_Decrypt
* argv[0]:Length in bytes of pbData
* argv[1]:Number of bytes to be decrypted by the cipher
* argv[2]:DRM_CIPHER_CONTEXT: NULL, NORMAL
*/
DRM_RESULT TestDRMCipher_Decrypt(long argc, char **argv)
{
    DRM_RESULT dr = DRM_SUCCESS;
    DRM_DWORD cbData = 0;
    int i,Datalength=0 ;
    DRM_BYTE* pbData = NULL;

    ChkArg(argc == 3);

    /* Get the encrypted data */
    if (argv[1]) 
    {
        Datalength = OEM_atoi(argv[1]);
        pbData = (DRM_BYTE*)OEM_malloc(Datalength);
        DX_VOS_FastMemCpy( pbData,rgbEncData,Datalength);
    }

    /* set content key length */
    if (argv[0]) {
        i = OEM_atoi(argv[0]);
        cbData = i;
    }

    /* set up cipher context */
    pContext = argv[2]? &oContext: NULL;

    ChkDR( DRM_CPHR_InitDecrypt( pContext, 
                                 (cbData < 15 ? pbData : pbData + cbData - 15),
                                 cbData ) );
    ChkDR(DRM_CPHR_Decrypt(pContext,cbData, pbData));

    /*compare the decrypted data with the clear content before encrypted */
    if (DX_VOS_MemCmp(rgbClearData, pbData, Datalength-1 ))
    {
        Log("ERROR", "CryptoTest.cpp: DRMCipher_Decrypt failed");
        dr = DRM_E_FAIL;
        goto ErrorExit;        
    }

ErrorExit:
    OEM_free(pbData);
    return dr;
}

/* 
*  TestShaInit - will initialize internal SHA data.  This must be called each time before a new SHA hash is computed.
*/
DRM_RESULT TestShaInit(long argc, char **argv)
{
    DRM_RESULT dr = DRM_SUCCESS;
    ChkArg(argv == 0);
    DRM_SHA_Init(&g_oShaContext);
ErrorExit:
    return dr;
}

/*
*  TestShaUpdate- updates the running hash context with new data to be hashed.  
*  argv[0]:index to the internal test buffer.
*  argv[2]:test buffer size
*/
DRM_RESULT TestShaUpdate(long argc, char **argv)
{
    DRM_RESULT dr = DRM_SUCCESS;
    DRM_DWORD i;

    ChkArg(argc == 2 && argv[0] && argv[1]);

    i = OEM_atol(argv[0]);
    ChkArg(i < MAX_DATA_COUNT);

    DRM_SHA_Update(g_pbRandomData[i], OEM_atol(argv[1]), &g_oShaContext);
 ErrorExit:
    return dr;
}

/*
* TestShaDone - signals that all data has been added to the hash 
* argv[0]: index to the internal buffer to store the hash result
*/
DRM_RESULT TestShaDone(long argc, char **argv)
{
    DRM_RESULT dr = DRM_SUCCESS;
     DRM_DWORD i;
    
    ChkArg(argc == 1 && argv[0]);

    i = OEM_atol(argv[0]);
    ChkArg(i < MAX_DATA_COUNT);
    DX_VOS_MemFree(g_pbRandomData[i]);
    ChkMem(g_pbRandomData[i] = (DRM_BYTE*)DX_VOS_MemMalloc(SHA_DIGEST_LEN));

    DRM_SHA_Finalize(&g_oShaContext, g_pbRandomData[i]);
ErrorExit:
    return dr;
}

/*
*   TestHMACInit - will initialize internal HMAC data.  
*    argv[0]:HMAC_CONTEXT: NULL, NORMAL
*    argv[1]:pointer to the pbKey to the hash
*    argv[2]:Length in bytes to add to the hash
*/

DRM_RESULT TestHMACInit(long argc, char **argv)
{
    DRM_RESULT dr = DRM_SUCCESS;
    DRM_BYTE* pbTestKey=NULL;
    DRM_UINT cbTestKey=0;
    int i,Datalength;
    DRM_BYTE rgbTempBuffer[SHA_BLOCK_SIZE];    /* This is 64 bytes.  */
    HMAC_CONTEXT *pHmacContext;

    pHmacContext = argv[0]? &gHmacContext : NULL;

    /* Pointer to an array of bytes with pbData to be added to the hash */
    if (argv[1]) 
    {
        Datalength =OEM_atoi(argv[1]);
        pbTestKey = (DRM_BYTE*)OEM_malloc(Datalength);
        OEM_GenRandomBytes(pbTestKey, Datalength );
    }

    /* set size of data to be hashed */
    if (argv[2]) {
        i = OEM_atoi(argv[2]);
        cbTestKey = i;
    }
    /* zero the mamory for compare */
    DX_VOS_MemSet( rgbTempBuffer, 0, SHA_BLOCK_SIZE );

    ChkDR(DRM_HMAC_Init( pHmacContext, pbTestKey, cbTestKey));

#ifndef DX_WMDRM_USE_CRYS
    /*verify the HMAC context after HMACInit */
    ChkArg(  !pHmacContext->shaContext.dwHighByteCount 
          && pHmacContext->shaContext.dwLowByteCount 
          && DX_VOS_MemCmp(rgbTempBuffer, pHmacContext->rgbBuffer, SHA_BLOCK_SIZE )
          && pHmacContext->shaContext.ABCDE[0]
          && pHmacContext->shaContext.ABCDE[1]
          && pHmacContext->shaContext.ABCDE[2]
          && pHmacContext->shaContext.ABCDE[3]
          && pHmacContext->shaContext.ABCDE[4]);
#endif

ErrorExit:
    OEM_free(pbTestKey);
    return dr;
}

/*
*   TestHMACAdd - updates the running hash context with new data to be hashed
*    argv[0]:HMAC_CONTEXT: NULL, NORMAL
*    argv[1]:pointer to the pbData to be added to the hash
*    argv[2]:Length in bytes to actually add to the hash
*/

DRM_RESULT TestHMACAdd(long argc, char **argv)
{
    DRM_RESULT dr = DRM_SUCCESS;
    DRM_BYTE* pbData=NULL;
    DRM_UINT cbData=0;
    DRM_UINT i,Datalength;
#ifndef DX_WMDRM_USE_CRYS
    DRM_BYTE rgbTempBuffer[SHA_BLOCK_SIZE];    /* This is 64 bytes.  */
#endif
    HMAC_CONTEXT oHmacContext={0};
    HMAC_CONTEXT *pHmacContext;

    pHmacContext = argv[0]? &gHmacContext: NULL;

    /* Pointer to an array of bytes with pbData to be added to the hash */
    if (argv[1]) 
    {
        Datalength = OEM_atoi(argv[1]);
        pbData = (DRM_BYTE*)OEM_malloc(Datalength);
        OEM_GenRandomBytes(pbData, Datalength );
    }

    /* set size of data to be hashed */
    if (argv[2]) {
        i = OEM_atoi(argv[2]);
        cbData = i;
    }
    /* zero the mamory for compare */

    if(pHmacContext)
        DX_VOS_FastMemCpy( &oHmacContext, pHmacContext, sizeof(HMAC_CONTEXT));

    ChkDR(DRM_HMAC_Update(pHmacContext,pbData,cbData));

    /*verify the HMAC context after HMACInit */
#ifndef DX_WMDRM_USE_CRYS
    ChkArg(  !pHmacContext->shaContext.dwHighByteCount 
          && (pHmacContext->shaContext.dwLowByteCount == oHmacContext.shaContext.dwLowByteCount + cbData)
          && DX_VOS_MemCmp(rgbTempBuffer, pHmacContext->rgbBuffer, SHA_BLOCK_SIZE ));
#endif
ErrorExit:
    OEM_free(pbData);
    return dr;
}


/*  TestHMACDone
*    argv[0]:HMAC_CONTEXT: NULL, NORMAL
*    argv[1]:An array of bytes to be filled with the resultant hash.
*    argv[2]:Length in bytes of cbKeyedHash Only a maximum of 20 bytes can be returned.
*/

DRM_RESULT TestHMACDone(long argc, char **argv)
{
    DRM_RESULT dr = DRM_SUCCESS; 
    int Datalength,i;
    DRM_BYTE* pbKeyedHash=NULL;
    DRM_UINT cbKeyedHash=0;
    HMAC_CONTEXT *pHmacContext;


    pHmacContext = argv[0]? &gHmacContext: NULL;
    if (argv[1]) 
    {
        Datalength = OEM_atoi(argv[1]);

        pbKeyedHash = (DRM_BYTE*)OEM_malloc(Datalength);
        OEM_GenRandomBytes(pbKeyedHash, Datalength );
    }

        /* set size of data to be hashed */
    if (argv[2]) {
        i = OEM_atoi(argv[2]);
        cbKeyedHash = i;
    }
    ChkDR( DRM_HMAC_Finalize( pHmacContext, pbKeyedHash, cbKeyedHash ) ); 

    if(!pbKeyedHash)
    {
        Log("ERROR", "CryptoTest.cpp: HMACDone failed");
        dr = DRM_E_FAIL;
        goto ErrorExit;        
    } 

ErrorExit:
    OEM_free(pbKeyedHash);
    return dr;
}

/* TestDRM_rc4_key
*  argv[0]: RC4_KEYSTRUCT: NORMAL, NULL
*  argv[1]: size of pRC4key in bytes
*  argv[2]: size of pbData being encrypt/decrypt in bytes
*/

DRM_RESULT TestDRM_rc4(long argc, char **argv)
{    
    DRM_RESULT dr = DRM_SUCCESS; 
    RC4_KEYSTRUCT *pRC4ks=NULL;
    RC4_KEYSTRUCT pTmpRC4ks={0};
    DRM_BYTE* pRC4key=NULL;
    DRM_BYTE* pbData=NULL;
    DRM_BYTE* pbTmpData=NULL;
    DRM_UINT  dwKeyLen=0,dwDataSize;
        
    pRC4ks = argv[0]? &rc4ks:NULL;

    if (argv[1]) {
        dwKeyLen = OEM_atoi(argv[1]);
        pRC4key = (DRM_BYTE*)OEM_malloc(dwKeyLen);
        OEM_GenRandomBytes(pRC4key, dwKeyLen );
    }

    if (argv[2]) 
    {
        dwDataSize = OEM_atoi(argv[2]);
        pbData = (DRM_BYTE*)OEM_malloc(dwDataSize);
        OEM_GenRandomBytes(pbData, dwDataSize );

        pbTmpData = (DRM_BYTE*)OEM_malloc(dwDataSize);
    }

    /*save the pRC4Key before calling the DRM_rc4_key */
    DX_VOS_FastMemCpy(&pTmpRC4ks,pRC4ks,sizeof(RC4_KEYSTRUCT));
    DRM_RC4_KeySetup( pRC4ks,dwKeyLen,pRC4key);

    /*compare the pRC4ks after calling the DRM_rc4_key, pRC4ks should not be the same*/
    if(!DX_VOS_MemCmp(&pTmpRC4ks,pRC4ks,sizeof(RC4_KEYSTRUCT)))
    {
        Log("ERROR", "CryptoTest.cpp: DRM_rc4_key failed");
        dr = DRM_E_FAIL;
        goto ErrorExit;        
    } 

    DX_VOS_FastMemCpy(pbTmpData,pbData,dwDataSize);

    /* The first call to DRM_rc4 encrypts pbData */
    DRM_RC4_Cipher( pRC4ks, dwDataSize, pbData );

    /*compare the pbData after calling the DRM_rc4, pbData should not be the same*/
    if(!DX_VOS_MemCmp(pbTmpData,pbData,dwDataSize))
    {
        Log("ERROR", "CryptoTest.cpp: DRM_rc4 failed");
        dr = DRM_E_FAIL;
        goto ErrorExit;        
    } 

    /*save the pRC4ks before calling the DRM_rc4_key */
    DX_VOS_FastMemCpy(&pTmpRC4ks,pRC4ks,sizeof(RC4_KEYSTRUCT));
    DRM_RC4_KeySetup( pRC4ks,dwKeyLen,pRC4key);
    
    /*compare the pRC4ks after calling the DRM_rc4_key, pRC4ks should not be the same*/
    if(!DX_VOS_MemCmp(&pTmpRC4ks,pRC4ks,sizeof(RC4_KEYSTRUCT)))
    {
        Log("ERROR", "CryptoTest.cpp: DRM_rc4_key failed");
        dr = DRM_E_FAIL;
        goto ErrorExit;        
    } 

    DRM_RC4_Cipher( pRC4ks, dwDataSize, pbData );

    /* compare the pbData with the one before being encrypted, it should be the same */
    if(DX_VOS_MemCmp(pbTmpData,pbData,dwDataSize))
    {
        Log("ERROR", "CryptoTest.cpp: DRM_rc4 failed");
        dr = DRM_E_FAIL;
        goto ErrorExit;        
    } 

ErrorExit:
    OEM_free(pRC4key);
    OEM_free(pbData);
    OEM_free(pbTmpData);
    return dr;
}

/* TestDeskey_Des
*  argv[0]: DESTable: NORMAL, NULL
*  argv[1]: size of pkey in bytes
*  argv[2]: size of pbIn in bytes to be encrypted 
*  argv[3]: size of pbOut in bytes to hold encrypted data 
*/

DRM_RESULT TestDeskey_Des(long argc, char **argv)
{    
    DRM_RESULT dr = DRM_SUCCESS; 
    DESTable  *pTable = NULL;
    DESTable pTmpTable={0};
    DRM_BYTE *pKey=NULL;
    DRM_BYTE *pbIn=NULL;
    DRM_BYTE *pbOut=NULL;
    DRM_BYTE *pbOrgIn=NULL;
    DRM_UINT  dwKeyLen=0,dwInDataSize,dwOutDataSize;
        
    pTable = argv[0]? &desTable:NULL;

    if (argv[1]) {
        dwKeyLen = OEM_atoi(argv[1]);
        pKey = (DRM_BYTE*)OEM_malloc(dwKeyLen);
        OEM_GenRandomBytes(pKey, dwKeyLen );    
    }
    if (argv[2]) 
    {
        dwInDataSize = OEM_atoi(argv[2]);
        pbIn = (DRM_BYTE*)OEM_malloc(dwInDataSize);
        OEM_GenRandomBytes(pbIn, dwInDataSize );

        pbOrgIn = (DRM_BYTE*)OEM_malloc(dwInDataSize);
        DX_VOS_FastMemCpy(pbOrgIn,pbIn,dwInDataSize);
    }
    if (argv[3]) 
    {
        dwOutDataSize = OEM_atoi(argv[3]);
        pbOut = (DRM_BYTE*)OEM_malloc(dwOutDataSize);
        OEM_GenRandomBytes(pbOut, dwOutDataSize );
    }

    DX_VOS_FastMemCpy(&pTmpTable,pTable,sizeof(DESTable));
    DRM_DES_KeySetup(pTable, pKey);  

    /*compare the pTable after calling the deskey, pTable should not be the same*/
    if(!DX_VOS_MemCmp(&pTmpTable,pTable,sizeof(DESTable)))
    {
        Log("ERROR", "CryptoTest.cpp: deskey failed");
        dr = DRM_E_FAIL;
        goto ErrorExit;        
    } 
    /* encrypt the pbIn, pbOut will hold the encrypted data */
    DRM_DES_Cipher( pbIn, pbOut, pTable, DES_ENCRYPT );
    
    /*compare the data after encrypted, pbOut should not be the same as pbOrgIn*/
    if(!DX_VOS_MemCmp(pbOrgIn,pbOut,dwOutDataSize))
    {
        Log("ERROR", "CryptoTest.cpp: des failed");
        dr = DRM_E_FAIL;
        goto ErrorExit;        
    } 
    
    /* Will decrypt the encrypted data */
    DRM_DES_KeySetup(pTable, pKey);  

    /* save the last pbOut as a new pbIn for next des call*/
    DX_VOS_FastMemCpy(pbIn,pbOut,dwInDataSize);
    DRM_DES_Cipher( pbIn, pbOut, pTable, DES_DECRYPT );

    /*compare the decrypted data with the original, pbOut should be the same as pbOrgIn*/
    if(DX_VOS_MemCmp(pbOrgIn,pbOut,DES_BLOCKLEN))
    {
        Log("ERROR", "CryptoTest.cpp: des failed");
        dr = DRM_E_FAIL;
        goto ErrorExit;        
    } 
    
ErrorExit:
    OEM_free(pKey);
    OEM_free(pbIn);
    OEM_free(pbOrgIn);
    OEM_free(pbOut);

    return dr;
}

/* TestCBC64WS4_asm
* argv[0]: cbData: length in bytes of input data
*/
DRM_RESULT TestCBC64WS4_asm(long argc, char **argv)
{
    DRM_RESULT dr;
    DRM_UINT   macLength4; /* mac length in four byte blocks */
    DRM_UINT   rc4key[ 2 ], *pLast8;    
/*    DRM_CBCKey      *pMackey = NULL;*/
    DRM_BYTE  *pbData = NULL;
    DRM_DWORD cbData=0;

    DRM_DWORD cbContentKey = DRMCIPHERKEYLEN;
    DRM_BYTE* pbContKey = NULL;
    
    ChkArg(argc == 1 && argv[0]);

    ChkMem(pbContKey = (DRM_BYTE*)OEM_malloc(cbContentKey));
    OEM_GenRandomBytes(pbContKey, cbContentKey);
    
    /* set up cipher context */
    pContext = &oContext;
    ChkDR(DRM_CPHR_Init(pContext, cbContentKey, pbContKey));

    cbData = OEM_atoi(argv[0]);

    ChkMem(pbData = (DRM_BYTE*)OEM_malloc(cbData));
    OEM_GenRandomBytes(pbData, cbData );

    ChkMem(pbOrgData = (DRM_BYTE*)OEM_malloc(cbData));
    DX_VOS_FastMemCpy(pbOrgData,pbData,cbData);
    ChkMem(pbEncData = (DRM_BYTE*)OEM_malloc(cbData));


    /* small packet case: MAC does not handle it */
    if ( cbData < 16 )
    {        
        for ( rc4key[0]=0; rc4key[0] < cbData; rc4key[0]++ )
        {
            pbData[rc4key[0]] ^= pContext->m_shaOut[rc4key[0]];
        }        
    }
    else
    {
        /*-----standard case: cbData >= 16----------------------------------------------------- */
        macLength4 = (cbData/8) * 2; /* making sure block number is even */
        pLast8 = (DRM_UINT *)(pbData + 4 * macLength4 - 8);

        /* run MAC over data */
        DRM_CBC_Mac( pbData, macLength4, (DRM_DWORD*)rc4key, &pContext->m_mackey);
 
        /* RC4 encrypt content */
        DRM_RC4_KeySetup( &pContext->m_rc4ks, 8, (DRM_BYTE *)rc4key );
        DRM_RC4_Cipher( &pContext->m_rc4ks, cbData, pbData );

        /* DES encrypt MAC and put it in the right place */
        rc4key[0] ^= pContext->m_desS1[0];
        rc4key[1] ^= pContext->m_desS1[1];
        DRM_DES_Cipher( (DRM_BYTE *)rc4key, (DRM_BYTE*)pLast8, &pContext->m_destable, DES_ENCRYPT );
        pLast8[0] ^= pContext->m_desS2[0];
        pLast8[1] ^= pContext->m_desS2[1];
    }
    /* save the pbData to pbEncData */
    DX_VOS_FastMemCpy(pbEncData,pbData,cbData);
    DX_VOS_FastMemCpy(&oContext,pContext,sizeof(DRM_CIPHER_CONTEXT));

ErrorExit:
    OEM_free(pbContKey);
    OEM_free(pbData);

    return dr;
}

/* TestInvCBC64WS4_asm
* argv[0]: cbData: length in bytes of input data
*/
DRM_RESULT TestInvCBC64WS4_asm(long argc, char **argv)
{
    DRM_RESULT dr = DRM_SUCCESS;
    DRM_UINT   macLength4; /* mac length in four byte blocks */
    DRM_UINT   rc4key[ 2 ], 
              *pLast8;
    DRM_BYTE  *pbData = NULL;
    DRM_DWORD  cbData=0;

    ChkArg(argc == 1 && argv[0]);

    /* set up cipher context */
    pContext = &oContext;

    pbData = pbEncData;

    cbData = OEM_atoi(argv[0]);

   /* small packet case: MAC does not handle it */
    if ( cbData < 16 )
    {
        for ( rc4key[0]=0; rc4key[0] < cbData; rc4key[0]++ )
        {
            pbData[rc4key[0]] ^= pContext->m_shaOut[rc4key[0]];
        }
    }
    else
    {
        /* making sure block number is even */
        macLength4 = (cbData / 8) * 2;
        pLast8 = (DRM_UINT *)(pbData + 4 * macLength4 - 8);

        /* DesS decrypt RC4 key */
        pLast8[0] ^= pContext->m_desS2[0];
        pLast8[1] ^= pContext->m_desS2[1];
        DRM_DES_Decrypt( (DRM_BYTE *)pLast8, (DRM_BYTE *)rc4key, &pContext->m_destable );
        rc4key[0] ^= pContext->m_desS1[0];
        rc4key[1] ^= pContext->m_desS1[1];

        /* RC4 decrypt content */
        DRM_RC4_KeySetup( &pContext->m_rc4ks, 8, (DRM_BYTE *)rc4key );
        DRM_RC4_Cipher( &pContext->m_rc4ks, cbData, pbData );

        /* take content of rc4key, reverse the mac and put it into the right position */
        pLast8[0] = rc4key[0]; 
        pLast8[1] = rc4key[1];
        DRM_CBC_InverseMac( pbData, macLength4, &(pContext->m_mackey), &(pContext->m_invmackey) );
    }
      
    /*compare the decrypted data with the original, pbOut should be the same as pbOrgIn*/
    if(DX_VOS_MemCmp(pbOrgData,pbData,sizeof(pbEncData)))
    {
        Log("ERROR", "CryptoTest.cpp: DRM_CBC_InverseMac failed");
        ChkDR(DRM_E_FAIL);
    } 
ErrorExit:
    return dr;
}

/* Get the default implementations for PreTestCase and PostTestCase */
DRM_RESULT CRY_PreTestCase(long lTCID, char *strTCName)
{
    int i;

    for (i = 0; i < MAX_DATA_COUNT; i++)
        g_pbRandomData[i] = NULL;

    bRandBufIn = (DRM_BYTE*)OEM_malloc(PK_ENC_PLAINTEXT_LEN );
    DX_VOS_MemSet( bRandBufIn, 0, PK_ENC_PLAINTEXT_LEN  );

    bCipher = (DRM_BYTE*)OEM_malloc(PK_ENC_CIPHERTEXT_LEN );
    DX_VOS_MemSet( bCipher, 0, PK_ENC_CIPHERTEXT_LEN  );

    bDecryptBufOut = (DRM_BYTE*)OEM_malloc(PK_ENC_PLAINTEXT_LEN );
    DX_VOS_MemSet( bDecryptBufOut, 0, PK_ENC_PLAINTEXT_LEN  );

    g_szKey = g_szPlainKey;
    prgbCryptoContext = (DRM_BYTE*)OEM_malloc(DRM_PKCRYPTO_CONTEXT_BUFFER_SIZE );
    DX_VOS_MemSet( prgbCryptoContext, 0, DRM_PKCRYPTO_CONTEXT_BUFFER_SIZE );

    pubKey = (PUBKEY*)OEM_malloc(PK_ENC_PUBLIC_KEY_LEN );
    DX_VOS_MemSet( pubKey, 0, PK_ENC_PUBLIC_KEY_LEN );

    privKey = (PRIVKEY*)OEM_malloc(PK_ENC_PRIVATE_KEY_LEN );
    DX_VOS_MemSet( privKey, 0, PK_ENC_PRIVATE_KEY_LEN );

    rgbEncData = (DRM_BYTE*)OEM_malloc(102400);
    rgbClearData = (DRM_BYTE*)OEM_malloc(102400);

    DX_VOS_MemSet( &m_mackey, 0, sizeof(DRM_CBCKey));
    DX_VOS_MemSet( &m_invmackey, 0, sizeof(DRM_CBCKey));

    return DRM_SUCCESS;
}

DRM_RESULT CRY_PostTestCase(long lTCID, char *strTCName)
{
    int i;

    for (i = 0; i < MAX_DATA_COUNT; i++) {
        DX_VOS_MemFree(g_pbRandomData[i]);
        g_pbRandomData[i] = NULL;
    }

    OEM_free(bRandBufIn);
    bRandBufIn = NULL;

    OEM_free(bCipher);
    bCipher = NULL;

    OEM_free(bDecryptBufOut);
    bDecryptBufOut = NULL;

    DX_VOS_MemSet( &oCrypto, 0, sizeof( DRM_CRYPTO_CONTEXT ) );

    OEM_free(prgbCryptoContext);
    prgbCryptoContext = NULL;

    OEM_free(pubKey);
    pubKey = NULL;
        
    OEM_free(privKey);
    privKey = NULL;

    OEM_free(pbEncData);
    pbEncData = NULL;

    OEM_free(pbOrgData);
    pbOrgData = NULL;

    OEM_free(rgbEncData);
    rgbEncData = NULL;

    OEM_free(rgbClearData);
    rgbClearData = NULL;

    return DRM_SUCCESS;
}

/*
IMPLEMENT_DEFAULT_WARPTEST

BEGIN_APIMAP(RefCryptoTest_ansi, "Crypto")
    API_ENTRY(TestPKEncrypt_PKDecrypt)
    API_ENTRY(TestPKSign_PKVerify)
    API_ENTRY(TestDRMCipher_Init)
    API_ENTRY(TestDRMCipher_Encrypt)
    API_ENTRY(TestDRMCipher_Decrypt)
    API_ENTRY(TestShaInit)
    API_ENTRY(TestShaUpdate)
    API_ENTRY(TestShaDone)
    API_ENTRY(TestHMACInit)
    API_ENTRY(TestHMACAdd)
    API_ENTRY(TestHMACDone)
    API_ENTRY(TestDRM_rc4)
    API_ENTRY(TestDeskey_Des)
    API_ENTRY(TestCBC64WS4_asm)
    API_ENTRY(TestInvCBC64WS4_asm)
    API_ENTRY(TestGenRandomData)
    API_ENTRY(TestMemComp)
END_APIMAP
*/
