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

#include <drmcommon.h>
#include <drmcipher.h>


DRM_RESULT DRM_API DRM_CPHR_Init(
    OUT       DRM_CIPHER_CONTEXT *pContext,
    IN        DRM_DWORD           cbContKey,
    IN  const DRM_BYTE           *pbContKey )
{
    DRM_RESULT  dr = DRM_E_LOGICERR;
    DRM_UINT    rnd_bytes[16];    
    DRM_UINT    rnd_bytes2[16];
    DRM_INT i=0;
    
    DRMCASSERT( SIZEOF( SHA_CONTEXT ) < SIZEOF( RC4_KEYSTRUCT ) );

    ChkArg (pbContKey  != NULL
        &&  cbContKey   > 0 
        &&  pContext   != NULL);

    ZEROMEM( pContext, SIZEOF( DRM_CIPHER_CONTEXT ) );
    /* hash to song key to 20 bytes length. first 12 bytes used as a MAC key, the following */
    /* 8 (7) bytes as a DES key */
    DRM_SHA_Init( (SHA_CONTEXT*)&pContext->m_rc4ks );
    DRM_SHA_Update( pbContKey, cbContKey, (SHA_CONTEXT*)&pContext->m_rc4ks );
    DRM_SHA_Finalize( (SHA_CONTEXT*)&pContext->m_rc4ks, pContext->m_shaOut );

    ZEROMEM( rnd_bytes2, SIZEOF( rnd_bytes ) );

    DRM_RC4_KeySetup( &pContext->m_rc4ks, 12, pContext->m_shaOut );
    DRM_RC4_Cipher( &pContext->m_rc4ks, SIZEOF( rnd_bytes2 ), (DRM_BYTE *)rnd_bytes2 );

    for( i=0; i<NO_OF(rnd_bytes2); i++)
    {
        BYTES_TO_DWORD( rnd_bytes[i], ((DRM_BYTE*)(rnd_bytes2 + i) ) );
    }
    
    pContext->m_mackey.a1 = rnd_bytes[0]  | 0x00000001;
    pContext->m_mackey.b1 = rnd_bytes[1]  | 0x00000001;
    pContext->m_mackey.c1 = rnd_bytes[2]  | 0x00000001;
    pContext->m_mackey.d1 = rnd_bytes[3]  | 0x00000001;
    pContext->m_mackey.e1 = rnd_bytes[4]  | 0x00000001;
    pContext->m_mackey.f1 = rnd_bytes[5]  | 0x00000001;
    pContext->m_mackey.a2 = rnd_bytes[6]  | 0x00000001;
    pContext->m_mackey.b2 = rnd_bytes[7]  | 0x00000001;
    pContext->m_mackey.c2 = rnd_bytes[8]  | 0x00000001;
    pContext->m_mackey.d2 = rnd_bytes[9]  | 0x00000001;
    pContext->m_mackey.e2 = rnd_bytes[10] | 0x00000001;
    pContext->m_mackey.f2 = rnd_bytes[11] | 0x00000001;

    pContext->m_invmackey.a1 = DRM_MAC_inv32( pContext->m_mackey.a1 );
    pContext->m_invmackey.a2 = DRM_MAC_inv32( pContext->m_mackey.a2 );
    pContext->m_invmackey.b1 = DRM_MAC_inv32( pContext->m_mackey.b1 );
    pContext->m_invmackey.b2 = DRM_MAC_inv32( pContext->m_mackey.b2 );
    pContext->m_invmackey.c1 = DRM_MAC_inv32( pContext->m_mackey.c1 );
    pContext->m_invmackey.c2 = DRM_MAC_inv32( pContext->m_mackey.c2 );
    pContext->m_invmackey.d1 = DRM_MAC_inv32( pContext->m_mackey.d1 );
    pContext->m_invmackey.d2 = DRM_MAC_inv32( pContext->m_mackey.d2 );
    pContext->m_invmackey.e1 = DRM_MAC_inv32( pContext->m_mackey.e1 );
    pContext->m_invmackey.e2 = DRM_MAC_inv32( pContext->m_mackey.e2 );

    pContext->m_desS1[0] = rnd_bytes[12]; 
    pContext->m_desS1[1] = rnd_bytes[13];
    pContext->m_desS2[0] = rnd_bytes[14]; 
    pContext->m_desS2[1] = rnd_bytes[15];

    DRM_DES_KeySetup( &pContext->m_destable, pContext->m_shaOut + __CB_DECL(12) );  
    pContext->m_fInited = TRUE;

    dr = DRM_SUCCESS;
ErrorExit:
    return dr;
}

/******************************************************************************
** 
** Function :   DRM_CPHR_InitDecrypt
** 
** Synopsis :   Initializes state for phased decryption
** 
** Arguments :  f_pcontextCipher - Cipher context initialized with DRM_CPHR_Init
**              f_pbLast15    - Last 15 bytes of ciphertext
**              f_cbData      - # of bytes in ENTIRE ciphertext              
** 
** Returns :    
** 
** Notes :      
** 
******************************************************************************/
DRM_RESULT DRM_API DRM_CPHR_InitDecrypt(
    IN  DRM_CIPHER_CONTEXT  *f_pcontextCipher,
    IN  DRM_BYTE            *f_pbLast15,
    IN  DRM_DWORD            f_cbData )
{
    DRM_DWORD  cbClear = 0;
    DRM_RESULT dr      = DRM_SUCCESS;    
        
    ChkArg( f_pcontextCipher != NULL
         && f_pbLast15       != NULL );

    if (!f_pcontextCipher->m_fInited)
    {
        ChkDR(DRM_E_CIPHER_NOTINITIALIZED);
    }
    
    /*
    **  Initialize data required for phased decryption
    */
    f_pcontextCipher->m_cbPacket       = f_cbData;
    f_pcontextCipher->m_iPacket        = 0;
    f_pcontextCipher->m_fDecryptInited = TRUE;
        
    if ( f_cbData < 16 )
    {        
        /* No need to do any other intialization  */
        goto ErrorExit;
    }

    /*
    **  Find Number of clear bytes
    */
    cbClear = f_cbData % 8;

    /*
    **  Extract the last 8 bytes before the clear content
    ** (To find these bytes start at the end of the last 15 bytes,
    **  subtract the number of clear bytes, then move 8 more bytes back,
    **  i.e. 15 - (cbClear + 8) ).
    */
    DRM_BYT_CopyBytes( (DRM_BYTE*)f_pcontextCipher->m_rguiLast8, 
                       0, 
                       f_pbLast15,
                       15 - ( cbClear + 8 ),
                       2 * SIZEOF(DRM_UINT) );    
    
    FIX_ENDIAN_DWORD( f_pcontextCipher->m_rguiLast8[0] );
    FIX_ENDIAN_DWORD( f_pcontextCipher->m_rguiLast8[1] );   

    /*    
    **  XOR the last 8 bytes(right before the clear content) with the DES key
    */    
    f_pcontextCipher->m_rguiLast8[0] ^= f_pcontextCipher->m_desS2[0];
    f_pcontextCipher->m_rguiLast8[1] ^= f_pcontextCipher->m_desS2[1];
    
    FIX_ENDIAN_DWORD( f_pcontextCipher->m_rguiLast8[0] );
    FIX_ENDIAN_DWORD( f_pcontextCipher->m_rguiLast8[1] );   
    
    /*
    **  Use the DES key to decrypt the 8 bytes. The result is the XOR'd RC4 
    **  content encryption key for the payload.
    */                
    DRM_DES_Decrypt( (DRM_BYTE *)f_pcontextCipher->m_rguiLast8, 
                     (DRM_BYTE *)f_pcontextCipher->m_rc4key, 
                     &f_pcontextCipher->m_destable );
    
    FIX_ENDIAN_DWORD( f_pcontextCipher->m_rguiLast8[0] );
    FIX_ENDIAN_DWORD( f_pcontextCipher->m_rguiLast8[1] );
    FIX_ENDIAN_DWORD( f_pcontextCipher->m_rc4key[0] );
    FIX_ENDIAN_DWORD( f_pcontextCipher->m_rc4key[1] );

    /*
    **  XOR the 8 bytes resulting from the DES decryption with the DES key. 
    **  The result is the RC4 content key for the packet
    */
    f_pcontextCipher->m_rc4key[0] ^= f_pcontextCipher->m_desS1[0];
    f_pcontextCipher->m_rc4key[1] ^= f_pcontextCipher->m_desS1[1];

    FIX_ENDIAN_DWORD( f_pcontextCipher->m_rc4key[0] );
    FIX_ENDIAN_DWORD( f_pcontextCipher->m_rc4key[1] );
    
    /*
    **  Setup the RC4 key to decrypt content
    */
    DRM_RC4_KeySetup( &f_pcontextCipher->m_rc4ks, 
                      8, 
                      (DRM_BYTE *)f_pcontextCipher->m_rc4key ); 
    
    FIX_ENDIAN_DWORD( f_pcontextCipher->m_rc4key[0] );
    FIX_ENDIAN_DWORD( f_pcontextCipher->m_rc4key[1] );

    /*
    **  Initialize cbc state
    */
    DRM_CBC64InitState( &(f_pcontextCipher->m_cbcstate) );
           
ErrorExit:
    return dr;
}
    
/******************************************************************************
** 
** Function :   DRM_CPHR_Decrypt
** 
** Synopsis :   Decrypts part of the ciphertext. This function must be called 
**              after DRM_CPHR_InitDecrypt has been successful
** 
** Arguments :  f_pcontextCipher - Cipher context initialized with DRM_CPHR_Init,
**                                 followed by DRM_CPHR_InitDecrypt
**              f_cbData  - # of bytes of data to be decrypted
**              f_pbData  - Pointer to encrypted buffer
** 
** Notes :      IMPORTANT: Phased decrypt should be called on segments of 
**              encrypted data sequentially, i.e. if encrypted data is divided 
**              four segments, the order of decryption should be segment 1,
**              followed by 2, 3, 4.
**              To remove complexity we do not handle the case where
**              a packet is split somewhere in the middle of the last 15 bytes.
**              The caller is responsible to ensure that the last 15 bytes never
**              straddle multiple calls to Decrypt.
******************************************************************************/
DRM_RESULT DRM_API DRM_CPHR_Decrypt(
    IN     DRM_CIPHER_CONTEXT *f_pcontextCipher,
    IN     DRM_DWORD           f_cbData, 
    IN OUT DRM_BYTE           *f_pbData )
{
    DRM_RESULT  dr = DRM_SUCCESS;
    DRM_DWORD   imac_start        = 0, 
                imac_end          = 0, 
                segment_end       = 0, 
                imac_in_seg_start = 0, 
                imac_in_seg_end   = 0;    
    
    ChkArg( f_pbData            != NULL
         && f_pcontextCipher    != NULL
         && (f_pcontextCipher->m_iPacket + f_cbData) <= f_pcontextCipher->m_cbPacket );

    if (!f_pcontextCipher->m_fInited)
    {
        ChkDR(DRM_E_CIPHER_NOTINITIALIZED);
    }

    if (!f_pcontextCipher->m_fDecryptInited)
    {
        ChkDR(DRM_E_DECRYPT_NOTINITIALIZED);
    }

    if (f_cbData == 0)
    {
        goto ErrorExit;
    }
    
    /*
    **  small packet case: MAC does not handle it
    */
    if ( f_pcontextCipher->m_cbPacket < 16 )
    {
        DRM_DWORD iData = 0;
        for ( iData = 0; iData < f_cbData; iData++)
        {
            DRM_BYTE bSHA = GET_BYTE( f_pcontextCipher->m_shaOut,
                                      iData + f_pcontextCipher->m_iPacket );
            DRM_BYTE bData = GET_BYTE( f_pbData, iData);

            PUT_BYTE( f_pbData, iData, bData ^ bSHA);
        }
        
        f_pcontextCipher->m_iPacket += f_cbData;
        goto ErrorExit;
    }   

    imac_end    = (f_pcontextCipher->m_cbPacket / 8) * 8;
    imac_start  = imac_end - 8;
    segment_end = f_pcontextCipher->m_iPacket + f_cbData;

    if ( segment_end > imac_start ) 
    {
        /* NOTE:  To remove complexity we do not handle the case where
           a packet is split somewhere in the middle of the last 15 bytes */
        DRMASSERT( segment_end == f_pcontextCipher->m_cbPacket );

        /* Set the last 8 bytes correctly */
        DRM_BYT_CopyBytes( f_pbData,
                           imac_start - f_pcontextCipher->m_iPacket,
                           (DRM_BYTE*)f_pcontextCipher->m_rguiLast8, 
                           0, 
                           2 * SIZEOF(DRM_UINT) );
    }

    /*
    **  RC4 decrypt the content
    */
    DRM_RC4_Cipher( &f_pcontextCipher->m_rc4ks, f_cbData, f_pbData );
    
    if ( f_pcontextCipher->m_iPacket < imac_start ) 
    {
        if ( f_pcontextCipher->m_iPacket + f_cbData >= imac_start ) 
        {
            DRM_UINT mac1 = 0;
            DRM_UINT mac2 = 0;
            DRM_UINT macInverse1 = 0;
            DRM_UINT macInverse2 = 0;

            /*
            **  First update MAC with data from this segment
            */            
            DRM_CBC64Update( &f_pcontextCipher->m_mackey, 
                         &f_pcontextCipher->m_cbcstate, 
                         imac_start - f_pcontextCipher->m_iPacket, 
                         f_pbData );
            
            /*
            **  Finalize MAC to decipher last 8 bytes of encrypted data 
            */
            mac1 = DRM_CBC64Finalize( &f_pcontextCipher->m_mackey, 
                                  &f_pcontextCipher->m_cbcstate, 
                                  &mac2 );
            macInverse2 = DRM_CBC64Invert( &f_pcontextCipher->m_mackey, 
                                       &f_pcontextCipher->m_invmackey, 
                                       mac1, 
                                       mac2, 
                                       f_pcontextCipher->m_rc4key[0], 
                                       f_pcontextCipher->m_rc4key[1], 
                                       &macInverse1 );
            f_pcontextCipher->m_rc4key[0] = macInverse1; 
            f_pcontextCipher->m_rc4key[1] = macInverse2;
        }
        else 
        {
            /*
            **  Update MAC with data from this segment
            */             
            DRM_CBC64Update( &f_pcontextCipher->m_mackey, 
                         &f_pcontextCipher->m_cbcstate, 
                         f_cbData, 
                         f_pbData );
        }
    }
         
    if ( f_pcontextCipher->m_iPacket < imac_end 
      && segment_end                 > imac_start ) 
    {
        /*
        **  Insert last 8 bytes of data deciphered
        */        
        DRM_BYTE  rgbMac[__CB_DECL(8)];
        DRM_DWORD iData = 0;
        
        DWORD_TO_BYTES( rgbMac,                f_pcontextCipher->m_rc4key[0] );
        DWORD_TO_BYTES( rgbMac + __CB_DECL(4), f_pcontextCipher->m_rc4key[1] );

        imac_in_seg_start = (imac_start >= f_pcontextCipher->m_iPacket) ? 
                                            imac_start : f_pcontextCipher->m_iPacket;
        imac_in_seg_end = (imac_end <= segment_end) ? imac_end:segment_end;
        
        for ( iData = imac_in_seg_start; iData < imac_in_seg_end; iData++ ) 
        {
            PUT_BYTE( f_pbData, 
                      iData - f_pcontextCipher->m_iPacket, 
                      GET_BYTE( rgbMac, iData - imac_start ) );
        }
    }    

    f_pcontextCipher->m_iPacket += f_cbData;
    if ( f_pcontextCipher->m_iPacket >= f_pcontextCipher->m_cbPacket )
    {
        f_pcontextCipher->m_fDecryptInited = FALSE;
    }


ErrorExit:
    return dr;
}

/*****************************************************************************/
DRM_RESULT DRM_API DRM_CPHR_Encrypt(
    IN     DRM_CIPHER_CONTEXT *pContext,
    IN     DRM_DWORD           cbData,
    IN OUT DRM_BYTE           *pbData )
{
    DRM_RESULT dr = DRM_SUCCESS;
    DRM_UINT   macLength4; /* mac length in four byte blocks */
    DRM_DWORD  rc4key    [2];
    DRM_DWORD  rguiLast8 [2];    

    ChkArg( pbData   != NULL 
         && pContext != NULL );
    
    if (!pContext->m_fInited)
    {
        ChkDR(DRM_E_CIPHER_NOTINITIALIZED);
    }
    
    /* small packet case: MAC does not handle it */
    if ( cbData < 16 )
    {        
        for ( rc4key[0]=0; rc4key[0] < cbData; rc4key[0]++ )
        {
            DRM_BYTE bTemp = GET_BYTE( pbData, rc4key[0] );
            bTemp ^= GET_BYTE( pContext->m_shaOut, rc4key[0] );
            PUT_BYTE( pbData, rc4key[0], bTemp );
        }        
    }
    else
    {
        /* making sure block number is even */
        macLength4 = (cbData / 8) * 2; 
        MEMCPY( rguiLast8, (pbData + __CB_DECL(4 * macLength4 - 8)), NO_OF(rguiLast8) * SIZEOF(DRM_UINT) );
        
        /* run MAC over data */
        DRM_CBC_Mac( pbData, macLength4, rc4key, &(pContext->m_mackey) );

        /* RC4 encrypt content */
        FIX_ENDIAN_DWORD( rc4key[0] );
        FIX_ENDIAN_DWORD( rc4key[1] );

        DRM_RC4_KeySetup( &pContext->m_rc4ks, 8, (DRM_BYTE *)rc4key );                
        DRM_RC4_Cipher(   &pContext->m_rc4ks, cbData, pbData );
        FIX_ENDIAN_DWORD( rc4key[0] );
        FIX_ENDIAN_DWORD( rc4key[1] );
        
        /* DES encrypt MAC and put it in the right place */
        rc4key[0] ^= pContext->m_desS1[0];
        rc4key[1] ^= pContext->m_desS1[1];
        
        FIX_ENDIAN_DWORD( rc4key[0] );
        FIX_ENDIAN_DWORD( rc4key[1] );
        DRM_DES_Encrypt( (DRM_BYTE *)rc4key, (DRM_BYTE*)rguiLast8, &pContext->m_destable);
        
        FIX_ENDIAN_DWORD( rguiLast8[0] );
        FIX_ENDIAN_DWORD( rguiLast8[1] );
        rguiLast8[0] ^= pContext->m_desS2[0];
        rguiLast8[1] ^= pContext->m_desS2[1];
        DWORD_TO_BYTES( (pbData + __CB_DECL(4 * macLength4 - 8)), rguiLast8[0] );
        DWORD_TO_BYTES( (pbData + __CB_DECL(4 * macLength4 - 4)), rguiLast8[1] );
    }
    dr = DRM_SUCCESS;
ErrorExit:
    return dr;
}

