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

#include <drmcommon.h>
#include <drmtypes.h>
#include <drmutilities.h>
#include <drmxmlbuilder.h>
#include <drmsha1.h>
#include <drmlicense.h>
#include <drmhdrbuilder.h>
#include <drmkeygenerator.h>


/*************************************************************************
** Macros
**************************************************************************/
#define LOCALSTACKSIZE  200
#define Tracked_N_ChkDR(hb,expr) {                         \
    if ( (hb)->dwItemsTracked == (hb)->dwItemsAdded ) \
    {                          \
        ChkDR((expr));         \
        ++(hb)->dwItemsAdded;  \
    }                          \
    ++(hb)->dwItemsTracked;    \
}


/*************************************************************************
** Local data structures
**************************************************************************/
typedef enum 
{
    eHBCreating=0x0,
    eHBInited, 
    eHBAddingData,
    eHBAddedData,
    eHBSigning,
    eHBCreated
} _EHBContextState;

/*
** Header Builder context stucture. This structure contains the following:
** - A Local stack
** - XML Builder context
** The XML Builder context is physically located at end of the buffer after init - for easier
** realloc a bigger size when necessary.
*/
typedef struct _tag_HBContext
{
    DRM_DWORD        wContextSize;
    _EHBContextState eState;
    DRM_DWORD        dwItemsAdded;     /* copied from old context when DRM_LA_ReallocRequest() is called */
    DRM_DWORD        dwItemsTracked;       /* 0 when DRM_LA_ReallocRequest() is called */
    DRM_STRING       szXMLString;
    DRM_BOOL         fKIDAdded;
    DRM_BOOL         fLAINFOAdded;
    DRM_BOOL         fChecksumAdded;
    DRM_BYTE         rgbLocalStack[__CB_DECL(LOCALSTACKSIZE)];     /* local stack buffer */
    _XMBContext     *pbXmlContext;
    _XMBContext      rgbXmlContext [1];
} _HBContext;


/*************************************************************************
** local constants
**************************************************************************/
/* default tag data */

static const DRM_WORD _cbXmlBuilderCtx=8192;    /* initial XML builder context size */

/*************************************************************************
** static functions
**************************************************************************/


/* 
**
*/
static DRM_VOID
_ChangeState(
    _HBContext *pHB,
    _EHBContextState eState)
{
    pHB->eState = eState;
    pHB->dwItemsAdded = pHB->dwItemsTracked = 0;
}


/*
** Add the following node/subnode to the request:
**
**  <SIGNATURE>
**    <HASHALGORITHM type=\"SHA\"></HASHALGORITHM>
**    <SIGNALGORITHM type=\"MSDRM\"></SIGNALGORITHM>
**    <VALUE>QrUgBBheJVzAOOw2kbo*1sqJkwft4oXKNBPv56PpUhho5fS6XK6sTg==</VALUE>
**  </SIGNATURE>
*/
static DRM_RESULT _SignTheHeader(
    _HBContext         *pHB,
    DRM_CRYPTO_CONTEXT *pCryptoContext,
    DRM_CONST_STRING   *pdstrPrivKey,
    DRM_STRING         *pdstrDataNodeFragment)
{
    DRM_RESULT dr=DRM_SUCCESS;

    ChkArg(pHB);

    /* open <SIGNATURE> node */
    Tracked_N_ChkDR(pHB, DRM_XMB_OpenNode(pHB->pbXmlContext, &g_dstrTagSignature));
        
    /* add <HASHALGORITHM> node */
    Tracked_N_ChkDR(pHB, DRM_XMB_OpenNode(pHB->pbXmlContext, &g_dstrTagHashAlg));
    Tracked_N_ChkDR(pHB, DRM_XMB_AddAttribute(pHB->pbXmlContext, &g_dstrAttributeType, &g_dstrSHA));
    Tracked_N_ChkDR(pHB, DRM_XMB_CloseCurrNode(pHB->pbXmlContext, NULL));

    /* add <SIGNALGORITHM> node */
    Tracked_N_ChkDR(pHB, DRM_XMB_OpenNode(pHB->pbXmlContext, &g_dstrTagSignAlg));
    Tracked_N_ChkDR(pHB, DRM_XMB_AddAttribute(pHB->pbXmlContext, &g_dstrAttributeType, &g_dstrMSDRM));
    Tracked_N_ChkDR(pHB, DRM_XMB_CloseCurrNode(pHB->pbXmlContext, NULL));

    /* add <VALUE> node */
    {
        DRM_BYTE    rgbSignature [__CB_DECL(PK_ENC_SIGNATURE_LEN)];
        DRM_DWORD   cchEncoded;
        PRIVKEY oPrivKey;
        DRM_DWORD dwPrivKey=SIZEOF(PRIVKEY);
        DRM_CONST_STRING dstrSignature = EMPTY_DRM_STRING;

        /* decode the signing key */
        ChkDR(DRM_B64_DecodeW(pdstrPrivKey, &dwPrivKey, (DRM_BYTE *)&oPrivKey, 0));

        /* sign the <DATA> node */        
        ChkDR(DRM_PK_Sign( pCryptoContext->rgbCryptoContext, 
                          &oPrivKey, 
                           PB_DSTR( pdstrDataNodeFragment ),
                           CB_DSTR( pdstrDataNodeFragment ),
                           rgbSignature));

        /* Base64 encode the signature */

        cchEncoded = CCH_BASE64_EQUIV( PK_ENC_SIGNATURE_LEN );
        if ( (cchEncoded + 1) * SIZEOF(DRM_WCHAR) > SIZEOF( pHB->rgbLocalStack ) )
        {
            ChkDR(DRM_E_OUTOFMEMORY);       /* LOCALSTACKSIZE too small */
        }
        ZEROMEM(pHB->rgbLocalStack, (cchEncoded+1) * SIZEOF(DRM_WCHAR) );
        ChkDR(DRM_B64_EncodeW( rgbSignature, 
                               PK_ENC_SIGNATURE_LEN, 
                  (DRM_WCHAR *)pHB->rgbLocalStack, 
                              &cchEncoded, 
                              0));
        dstrSignature.pwszString = (DRM_WCHAR *)pHB->rgbLocalStack;
        dstrSignature.cchString = cchEncoded;

        Tracked_N_ChkDR(pHB, DRM_XMB_OpenNode(pHB->pbXmlContext, &g_dstrTagValue));
        Tracked_N_ChkDR(pHB, DRM_XMB_AddData(pHB->pbXmlContext, &dstrSignature));
        Tracked_N_ChkDR(pHB, DRM_XMB_CloseCurrNode(pHB->pbXmlContext, NULL));
    }

    /* close SIGNATURE node */
    Tracked_N_ChkDR(pHB, DRM_XMB_CloseCurrNode(pHB->pbXmlContext, NULL));

ErrorExit:
    return dr;
}



DRM_RESULT DRM_API DRM_HB_CreateHeader(
    IN     DRM_BYTE  *pbHBContext,
    IN OUT DRM_DWORD *pcbHBContext)
{
    DRM_RESULT  dr        = DRM_SUCCESS;
    _HBContext *pHB       = NULL;
    DRM_DWORD   cbContext = 0;
    DRM_DWORD   cbTotal   = 0;
    DRM_DWORD   cbAligned = 0;
    DRM_DWORD   cbAdjust  = 0;

    ChkArg(pcbHBContext);

    /* check given context size */
    cbContext = SIZEOF(_HBContext) - 1 + _cbXmlBuilderCtx;
    if ( pbHBContext == NULL || *pcbHBContext < cbContext )
    {
        /* minimum context size to start with */
        *pcbHBContext= cbContext;
        dr = DRM_E_BUFFERTOOSMALL;
        goto ErrorExit;
    }

    /* Initialize Context 
    ** note: pbXmlContext is physically located at end of pbBlackboxContext.
    */
    ZEROMEM(pbHBContext, *pcbHBContext);
    pHB = (_HBContext*)pbHBContext;
    pHB->wContextSize = (DRM_WORD)*pcbHBContext;

    /* init initial state */
    _ChangeState(pHB, eHBCreating);

    cbTotal = *pcbHBContext - SIZEOF(_HBContext) + 1;

    ChkDR(DRM_UTL_EnsureDataAlignment((DRM_BYTE*)pHB->rgbXmlContext, 
                                       cbTotal, 
                                       (DRM_BYTE**)&pHB->pbXmlContext, 
                                       &cbAligned, 
                                       SIZEOF(DRM_DWORD), 
                                       &cbAdjust));

    pHB->wContextSize -= cbAdjust;
    
    /* Init XML Builder 
    ** Add document root tag: <WRMHEADER version=2.0.0.0"> 
    */
    ChkDR(DRM_XMB_CreateDocument(cbAligned, 
                                 pHB->pbXmlContext, 
                                &g_dstrTagWrmHeader));
                                 
    ChkDR(DRM_XMB_AddAttribute(pHB->pbXmlContext, &g_dstrAttributeVersion, &g_dstrAttributeVersionValue));

    /* open <DATA> node to the root node  */
    ChkDR(DRM_XMB_OpenNode(pHB->pbXmlContext, &g_dstrTagData));

    /* change state */
    _ChangeState(pHB, eHBInited);

ErrorExit:
    return dr;
}


/*
**
*/
DRM_RESULT DRM_API DRM_HB_ReallocHeader(
    IN DRM_BYTE *pbOldHBContext,
    IN DRM_DWORD cbNewHBContext,
    IN DRM_BYTE *pbNewHBContext)
{
    DRM_RESULT  dr     = DRM_SUCCESS;
    _HBContext *pOldHB = (_HBContext*)pbOldHBContext;
    _HBContext *pNewHB = (_HBContext*)pbNewHBContext;

    ChkArg(pbOldHBContext && pbNewHBContext  && pOldHB->eState!=eHBCreating
        && cbNewHBContext >= SIZEOF(_HBContext)   /* Prevent potential overflow */
        && cbNewHBContext > pOldHB->wContextSize);

    /* copy old context to new buffer */
    ZEROMEM(pbNewHBContext,  cbNewHBContext);
    pNewHB->wContextSize   = cbNewHBContext;
    pNewHB->dwItemsAdded   = pOldHB->dwItemsAdded;
    pNewHB->eState         = pOldHB->eState;
    pNewHB->fKIDAdded      = pOldHB->fKIDAdded;
    pNewHB->fLAINFOAdded   = pOldHB->fLAINFOAdded;
    pNewHB->fChecksumAdded = pOldHB->fChecksumAdded;
    
    ChkDR(DRM_XMB_ReallocDocument(pOldHB->pbXmlContext, 
        cbNewHBContext-SIZEOF(_HBContext)+1, pNewHB->pbXmlContext));

ErrorExit:
    return dr;
}


/*
**
*/
DRM_RESULT DRM_API DRM_HB_GetContextSize(
    IN  DRM_BYTE  *pbHBContext,
    OUT DRM_DWORD *pcbHBContext)
{
    DRM_RESULT dr=DRM_SUCCESS;
    _HBContext *pHB=(_HBContext*)pbHBContext;
    
    ChkArg(pbHBContext && pcbHBContext);  /* && pHB->eState!=eHBCreating); */

    /* current context size */
    *pcbHBContext = pHB->wContextSize;

ErrorExit:
    return dr;
}


/*
** 
*/
DRM_RESULT DRM_API DRM_HB_SignHeader(
    IN  DRM_BYTE           *pbHBContext,
    IN  DRM_CRYPTO_CONTEXT *pCryptoContext,
    IN  DRM_CONST_STRING   *pdstrPrivKey,      /* base64 encoded */
    OUT DRM_BYTE          **ppbHeaderXML,
    OUT DRM_DWORD          *pcbHeaderXML)
{
    DRM_RESULT dr=DRM_SUCCESS;
    _HBContext *pHB=(_HBContext*)pbHBContext;    
    DRM_STRING dstrDataNodeFragment;

    ChkArg(pbHBContext && pcbHeaderXML && ppbHeaderXML);
    ChkDRMString(pdstrPrivKey);
    if (pHB->eState!=eHBAddedData && pHB->eState!=eHBSigning)
    {
        if (!pHB->fKIDAdded)
        {
            ChkDR(DRM_E_NOKIDINHEADER);
        }
        else if (!pHB->fLAINFOAdded)
        {
            ChkDR(DRM_E_NOLAINFOINHEADER);
        }
        else if (!pHB->fChecksumAdded)
        {
            ChkDR(DRM_E_NOCHECKSUMINHEADER);
        }
        else
        {
            ChkDR(DRM_E_CH_INVALID_HEADER);
        }
    }

    if ( pHB->eState!= eHBSigning )
    {
        _ChangeState(pHB, eHBSigning);
    }

    /* close <DATA> node and get the XML fragment of it */
    Tracked_N_ChkDR(pHB, DRM_XMB_CloseCurrNode(pHB->pbXmlContext, &dstrDataNodeFragment));
    /* Sign the <DATA> node */
    ChkDR(_SignTheHeader(pHB, pCryptoContext, pdstrPrivKey, &dstrDataNodeFragment));
    Tracked_N_ChkDR(pHB, DRM_XMB_CloseDocument(pHB->pbXmlContext, &pHB->szXMLString));

    *ppbHeaderXML = PB_DSTR( &pHB->szXMLString );
    *pcbHeaderXML = CB_DSTR( &pHB->szXMLString );
    
    /* make sure the context will not be used to call the API again */
    _ChangeState(pHB, eHBCreated);
    
ErrorExit:
    return dr;
}


/*
**
*/
DRM_RESULT DRM_API DRM_HB_SetAttribute(
    IN       DRM_BYTE         *pbHBContext,
    IN const DRM_CONST_STRING *pdstrAttrName,
    IN const DRM_CONST_STRING *pdstrAttrValue)
{
    DRM_RESULT dr=DRM_SUCCESS;
    _HBContext *pHB=(_HBContext*)pbHBContext;
    
    ChkArg(pbHBContext);
    ChkDRMString(pdstrAttrName);
    ChkDRMString(pdstrAttrValue);
    if ( pHB->eState != eHBInited &&
        pHB->eState != eHBAddingData&&
        pHB->eState != eHBAddedData )
    {
        ChkArg(FALSE);
    }

    /* Add <ACTION> node */
    if ( pHB->eState!=eHBAddingData )
    {
        _ChangeState(pHB, eHBAddingData);
    }

    Tracked_N_ChkDR(pHB, DRM_XMB_OpenNode(pHB->pbXmlContext, pdstrAttrName));
    Tracked_N_ChkDR(pHB, DRM_XMB_AddData(pHB->pbXmlContext, pdstrAttrValue));
    Tracked_N_ChkDR(pHB, DRM_XMB_CloseCurrNode(pHB->pbXmlContext, NULL));

    _ChangeState(pHB, eHBAddedData);
    
ErrorExit:
    return dr;
}

/*
**
*/
DRM_RESULT DRM_API DRM_HB_SetKeyID(
    IN       DRM_BYTE         *pbHBContext,
    IN const DRM_CONST_STRING *pdstrKID)
{
    DRM_RESULT dr=DRM_SUCCESS;
    _HBContext *pHB=(_HBContext*)pbHBContext;
    
    ChkArg(pbHBContext);
    ChkDRMString(pdstrKID);
    ChkArg ( pHB->eState == eHBInited || pHB->eState == eHBAddingData ||
        pHB->eState == eHBAddedData );

    if ( pHB->fKIDAdded )
    {
        ChkDR(DRM_E_DUPLICATEDHEADERATTRIBUTE);
    }

    /* set the attribute */
    ChkDR(DRM_HB_SetAttribute(pbHBContext, &g_dstrTagKID, pdstrKID));
    pHB->fKIDAdded = TRUE;

ErrorExit:
    return dr;
}

/*
**
*/
DRM_RESULT DRM_API DRM_HB_SetLicAcqUrl(
    IN       DRM_BYTE         *pbHBContext,
    IN const DRM_CONST_STRING *pdstrURL)
{
    DRM_RESULT dr=DRM_SUCCESS;
    _HBContext *pHB=(_HBContext*)pbHBContext;
    
    ChkArg(pbHBContext);
    ChkDRMString(pdstrURL);
    ChkArg ( pHB->eState == eHBInited || pHB->eState == eHBAddingData ||
        pHB->eState == eHBAddedData );

    if ( pHB->fLAINFOAdded )
    {
        ChkDR(DRM_E_DUPLICATEDHEADERATTRIBUTE);
    }
    ChkDR(DRM_HB_SetAttribute(pbHBContext, &g_dstrTagLAINFO, pdstrURL));
    pHB->fLAINFOAdded = TRUE;
    
ErrorExit:
    return dr;
}


/*
** Perform checksum algo to the given content key and set the check to header
*/
DRM_RESULT DRM_API DRM_HB_SetChecksum(
    IN       DRM_BYTE         *pbHBContext,
    IN const DRM_CONST_STRING *pdstrContentKey)
{
    DRM_RESULT dr = DRM_SUCCESS;
    _HBContext      *pHB        = (_HBContext*)pbHBContext;
    DRM_DWORD        cchEncoded = 0;
    DRM_INT          lCount     = 0;
    SHA_CONTEXT      shaVal     = { 0 };
    DRM_BYTE         res           [__CB_DECL(SHA_DIGEST_LEN + 1)]     = { 0 };
    DRM_BYTE         rgbKeyContent [__CB_DECL(DRM_CONTENT_KEY_LENGTH)] = { 0 };
    DRM_DWORD        cbContentKey = DRM_CONTENT_KEY_LENGTH;
    DRM_CONST_STRING dstrChecksum = EMPTY_DRM_STRING;
    
    ChkArg(pbHBContext);
    ChkDRMString(pdstrContentKey);
    ChkArg ( pHB->eState == eHBInited || pHB->eState == eHBAddingData ||
        pHB->eState == eHBAddedData );

    if ( pHB->fChecksumAdded )
    {
        ChkDR(DRM_E_DUPLICATEDHEADERATTRIBUTE);
    }

    ChkDR(DRM_B64_DecodeW(pdstrContentKey, &cbContentKey, rgbKeyContent, 0));


    /* Run SHA on key SHA_ITERATIONS times. */
    /* First copy key into res. */
    
    DRMCASSERT ( DRM_CONTENT_KEY_LENGTH < SHA_DIGEST_LEN+1 ); /* Make sure we don't exceed buffer capacity in res. */    
    ZEROMEM(res, SHA_DIGEST_LEN+1);
    MEMCPY (res, rgbKeyContent, cbContentKey);
    for (lCount = 0; lCount < SHA_ITERATIONS; lCount++)
    {
        DRM_SHA_Init(&shaVal);
        DRM_SHA_Update(res, SHA_DIGEST_LEN+1, &shaVal);
        DRM_SHA_Finalize(&shaVal, res);
    }

    /* take first CHECKSUM_LENGTH bytes of res and use it to encode. */
    cchEncoded = CCH_BASE64_EQUIV( CHECKSUM_LENGTH );
    if ( (cchEncoded * SIZEOF(DRM_WCHAR)) >= LOCALSTACKSIZE )
    {
        ChkDR(DRM_E_OUTOFMEMORY);       /* LOCALSTACKSIZE too small */
    }
    ZEROMEM(pHB->rgbLocalStack, cchEncoded * SIZEOF(DRM_WCHAR));
    ChkDR(DRM_B64_EncodeW(res, CHECKSUM_LENGTH, (DRM_WCHAR*)pHB->rgbLocalStack, &cchEncoded, 0));

    /* set the attribute */   
    dstrChecksum.pwszString = (DRM_WCHAR*)pHB->rgbLocalStack;
    dstrChecksum.cchString = cchEncoded;
    ChkDR(DRM_HB_SetAttribute(pbHBContext, &g_dstrTagChecksum, &dstrChecksum));
    pHB->fChecksumAdded = TRUE;

ErrorExit:
    return dr;
}


