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

#include <drmcommon.h>
#include <drmutilities.h>
#include <drmcrt.h>
#include <drmcontextsizes.h>
#include <drmsha1.h>
#include <drmkeygenerator.h>
#include <oemimpl.h>


DRM_RESULT DRM_API DRM_KG_GenerateKeyseed(
       OUT DRM_WCHAR *pwszKeyseed,
    IN OUT DRM_DWORD *pcchKeyseed)
{
    DRM_RESULT dr=DRM_SUCCESS;
    DRM_BYTE rgbKeyseed [__CB_DECL(DRM_KEYSEED_LEN)];
    DRM_INT i=0;

    ChkArg(pcchKeyseed);
    if ( *pcchKeyseed<(DRM_KEYSEED_LEN+1) || pwszKeyseed==NULL )
    {
        dr = DRM_E_BUFFERTOOSMALL;
        *pcchKeyseed = DRM_KEYSEED_LEN+1;
        goto ErrorExit;
    }
    *pcchKeyseed = DRM_KEYSEED_LEN+1;
    ChkDR(OEM_GenRandomBytes(rgbKeyseed, DRM_KEYSEED_LEN));

    ZEROMEM((DRM_BYTE*)pwszKeyseed, *pcchKeyseed * SIZEOF(DRM_WCHAR));
    for (i = 0; i < DRM_KEYSEED_LEN; i++)
    {
        /* convert each byte into an alpha numberic character (62 available) */
        DRM_BYTE val = (GET_BYTE(rgbKeyseed, i)) % 62;

        if (val < 10)
        {
            val += NATIVE_WCHAR(g_wch0);
        }
        else if (val < 36)
        {
            val += NATIVE_WCHAR(g_wcha) - 10;
        }
        else
        {
            val += NATIVE_WCHAR(g_wchA) - 36;
        }
        pwszKeyseed[i] = WCHAR_CAST(val);
    }

ErrorExit:

    return dr;
}

DRM_RESULT DRM_API DRM_KG_GenerateKeyID(
       OUT DRM_WCHAR *pwszKeyID,           /* base64 encoded string */
    IN OUT DRM_DWORD *pcchKeyID)
{
    DRM_RESULT dr=DRM_SUCCESS;
    DRM_GUID keyGuid;
    DRM_DWORD cchEncoded = CCH_BASE64_EQUIV(SIZEOF(DRM_GUID));

    ChkArg(pcchKeyID);

    if ( *pcchKeyID<( cchEncoded+1) 
      || pwszKeyID==NULL )
    {
        *pcchKeyID=cchEncoded+1;
        dr =DRM_E_BUFFERTOOSMALL;
        goto ErrorExit;
    }
    
    /*
    ** generate random bytes for KeyID (used to generate a GUID)
    */
    ChkDR(OEM_GenRandomBytes((DRM_BYTE*)&keyGuid, SIZEOF(DRM_GUID)));

    /* base64 encode */    
    ZEROMEM((DRM_BYTE*)pwszKeyID, *pcchKeyID * SIZEOF(DRM_WCHAR));
    ChkDR(DRM_B64_EncodeW((DRM_BYTE *)&keyGuid, SIZEOF(DRM_GUID), pwszKeyID, &cchEncoded, 0));
    *pcchKeyID=cchEncoded+1;

ErrorExit:
    return dr;
}


DRM_RESULT DRM_API DRM_KG_GenerateContentKey(
    IN const    DRM_WCHAR *pwszKeySeed,          /* base64 encoded string */
    IN          DRM_DWORD  cchKeySeed, 
    IN const    DRM_WCHAR *pwszKeyID,            /* base64 encoded string */
    IN          DRM_DWORD  cchKeyID, 
       OUT      DRM_WCHAR *pwszContentKey,       /* base64 encoded string */
    IN OUT      DRM_DWORD *pcchContentKey)
{
    DRM_RESULT dr=DRM_SUCCESS;
    SHA_CONTEXT  shaVal;    
    DRM_BYTE rgbKeyseed[__CB_DECL(DRM_KEYSEED_LEN)];
    DRM_BYTE res       [__CB_DECL(SHA_DIGEST_LEN + 1)];
    DRM_BYTE res1      [__CB_DECL(SHA_DIGEST_LEN + 1)];
    DRM_BYTE res2      [__CB_DECL(SHA_DIGEST_LEN + 1)];
    DRM_DWORD cchb64 = CCH_BASE64_EQUIV(DRM_CONTENT_KEY_LENGTH);
    DRM_UINT  i      = 0;    

    ChkArg( pwszKeySeed    != NULL
         && cchKeySeed     == DRM_KEYSEED_LEN 
         && pwszKeyID      != NULL
         && cchKeyID       != 0
         && pcchContentKey != NULL );

    /*
    **  We should actually limit this to just 24 bytes; but for historical
    **  reasons, we have some providers using slightly longer KIDs. Resuing the
    **  keyseed buffer allows for slightly longer KIDs while still checking for 
    **  KIDs that are too large
    */
    ChkArg( cchKeyID       < SIZEOF(rgbKeyseed) );
         
    /* check buffer length */
    if ( *pcchContentKey < cchb64+1 
      || pwszContentKey == NULL )
    {
        *pcchContentKey=cchb64+1;
        dr=DRM_E_BUFFERTOOSMALL;
        goto ErrorExit;
    }
    
    for( i=0; i<cchKeySeed; i++ )
    {
        PUT_BYTE(rgbKeyseed, i,  (DRM_BYTE) NATIVE_WCHAR (pwszKeySeed [i]));
    }

    /* init buffer */
    ZEROMEM((DRM_BYTE*)pwszContentKey, *pcchContentKey*SIZEOF(DRM_WCHAR));

    /* generate key */
    DRM_SHA_Init(&shaVal);
    DRM_SHA_Update(rgbKeyseed, DRM_KEYSEED_LEN, &shaVal);
    DRM_SHA_Finalize(&shaVal, res);

    for( i=0; i<cchKeyID; i++ )
    {
        PUT_BYTE(rgbKeyseed, i, (DRM_BYTE) NATIVE_WCHAR (pwszKeyID [i]));
    }

    DRM_SHA_Init(&shaVal);
    DRM_SHA_Update(res, 6, &shaVal);  /* Copy first 48 bits of res */
    DRM_SHA_Update(rgbKeyseed, cchKeyID, &shaVal);
    DRM_SHA_Finalize(&shaVal, res1);

    DRM_SHA_Init(&shaVal);
    DRM_SHA_UpdateOffset(res, 6, 6, &shaVal);  /* Copy 2nd set of 48 bits of res */
    DRM_SHA_Update(rgbKeyseed, cchKeyID, &shaVal);
    DRM_SHA_Finalize(&shaVal, res2);

    DRM_SHA_Init(&shaVal);
    DRM_SHA_UpdateOffset(res, 12, 6, &shaVal);  /* Copy third set of 48 bits of res */
    DRM_SHA_Update(rgbKeyseed, cchKeyID, &shaVal);
    DRM_SHA_Finalize(&shaVal, res);

    /* Reuse res to hold the final result. */
    for (i = 0; i < SHA_DIGEST_LEN; i++)
    {
        DRM_BYTE b = GET_BYTE(res, i);
        b ^= GET_BYTE(res1, i) ^ GET_BYTE(res2, i);
        PUT_BYTE(res, i, b);
    }
    
    ChkDR(DRM_B64_EncodeW(res, DRM_CONTENT_KEY_LENGTH, pwszContentKey, &cchb64, 0));    
    pwszContentKey[cchb64] = g_wchNull;
    *pcchContentKey=cchb64+1;

ErrorExit:
    return dr;
}


DRM_RESULT DRM_API DRM_KG_GenerateSigningKeys(
    IN     DRM_CRYPTO_CONTEXT *pCrypto,
       OUT DRM_WCHAR          *pwszPrivKey,         /* base64 encoded string */
    IN OUT DRM_DWORD          *pcchPrivKey,
       OUT DRM_WCHAR          *pwszPubKey,          /* base64 encoded string */
    IN OUT DRM_DWORD          *pcchPubKey)
{
    DRM_RESULT dr=DRM_SUCCESS;
    PUBKEY pub;
    PRIVKEY priv;
    DRM_DWORD cchPubEncoded  = CCH_BASE64_EQUIV(PK_ENC_PUBLIC_KEY_LEN);
    DRM_DWORD cchPrivEncoded = CCH_BASE64_EQUIV(PK_ENC_PRIVATE_KEY_LEN);

    if ( *pcchPrivKey <  cchPrivEncoded + 1 
      || *pcchPubKey  <  cchPubEncoded  + 1 
      ||  pwszPrivKey == NULL 
      ||  pwszPubKey  == NULL )
    {
        *pcchPrivKey = cchPrivEncoded + 1;
        *pcchPubKey  = cchPubEncoded  + 1;
        dr = DRM_E_BUFFERTOOSMALL;
        goto ErrorExit;
    }

    ZEROMEM((DRM_BYTE*) pwszPrivKey, *pcchPrivKey * SIZEOF(DRM_WCHAR));
    ZEROMEM((DRM_BYTE*) pwszPubKey,  *pcchPubKey  * SIZEOF(DRM_WCHAR));

    /* generate the key pair */
    ChkDR(DRM_PK_GenKeyPair( pCrypto->rgbCryptoContext, &pub, &priv ) );

    /* base64 encode */
    ChkDR(DRM_B64_EncodeW((DRM_BYTE *)&pub,  PK_ENC_PUBLIC_KEY_LEN,  pwszPubKey,  pcchPubKey,  0));    
    ChkDR(DRM_B64_EncodeW((DRM_BYTE *)&priv, PK_ENC_PRIVATE_KEY_LEN, pwszPrivKey, pcchPrivKey, 0));
    *pcchPubKey  += 1;
    *pcchPrivKey += 1;

ErrorExit:
    return dr;
}



