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

#include "blackboxtest.h"

static DRM_RESULT ArgToBindingInfo(char *arg, DRM_BINDING_INFO **ppB)
{
	int i;
	if (arg) {
		i = OEM_atoi(arg);
		if (i < 0 || i >= MAX_GLOBAL_BINDING_BLOB)
			return DRM_E_INVALIDARG;
		*ppB = &(g_binding[i]);
	} else
		*ppB = NULL;
	return DRM_SUCCESS;
}

/* Calculate the checksum */
static void SetCheckSum(DRM_BYTE *pData, DRM_DWORD cbData, DRM_BYTE pbCheckSum[CHECKSUM_LENGTH])
{
    DRM_BYTE res[SHA_DIGEST_LEN+1];
    SHA_CONTEXT shaVal;
    DRM_LONG lCount;

    /* Run SHA on key SHA_ITERATIONS times. */
    /* First copy key into res. */
    if ( cbData > SHA_DIGEST_LEN+1 )
        cbData = SHA_DIGEST_LEN+1; /* Make sure we don't exceed buffer capacity in res. */

    DX_VOS_MemSet( res, 0, sizeof(res) ); /* Init to 0. */
    DX_VOS_FastMemCpy( res, pData, cbData );   /* Copy key. */

    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 );
    }
    DX_VOS_FastMemCpy(pbCheckSum, res, CHECKSUM_LENGTH);
}

/* Return the blackbox pubkey to encrypt the content key. */
static DRM_RESULT GetBBPubKey(PUBKEY* pPubKey)
{
	DX_VOS_FastMemCpy(pPubKey, &(g_BBContext.cachedCertValues.m_BBCompatibilityCert.pk.pk), sizeof(PUBKEY));
    return DRM_SUCCESS;
}

/*
	Create internal BindingInfo structure for other api
	argv[0]: index to the global BindingInfo to create
	argv[1]: modifications to the BindingInfo:
		BAD_CHECKSUM, BAD_SIGNATURE, BAD_CONTENTKEY, BAD_LSPUBKEY, WRONG_LSPUBKEY, BAD_BBPUBKEY, KEY_LENGTH #
*/
DRM_RESULT Test_CreateBindingInfo(long argc, char **argv)
{
	DRM_RESULT dr;
	DRM_BINDING_INFO *pBindingInfo;
	DRM_BYTE keyLength, ClearContentKey[PK_ENC_PLAINTEXT_LEN];
	PRIVKEY lsPrvKey; /* license server private key to sign the content key */

	ChkArg(argc == 2 && argv[1]);
	ChkDR(ArgToBindingInfo(argv[0], &pBindingInfo));
	ChkArg(pBindingInfo);
	
	DX_VOS_MemSet(pBindingInfo, 0, sizeof(DRM_BINDING_INFO));
	
	/* Generate the random content key. The first byte is the key length. 
		The actual key starts from the second byte. */
	ChkDR(OEM_GenRandomBytes(ClearContentKey, PK_ENC_PLAINTEXT_LEN));
	
	/* Set the content key length */
	if (!DX_VOS_StrNCmp(argv[1], "KEY_LENGTH", 10))
		keyLength = (DRM_BYTE)OEM_atoi(argv[1] + 10);
	else
		keyLength = (DRM_BYTE)DRMCIPHERKEYLEN; /* drmcipher limits the key length to DRMCIPHERKEYLEN */
	ClearContentKey[0] = keyLength;

	/* Set the algorithm to ECC for now. Need to test chained licenses later using DES */
	pBindingInfo->m_dwAlgorithm = eMSDRM_PK;

	/* Copy the blackbox public key to do the encryption. This is the pubkey sent up to the license server. */
	ChkDR(GetBBPubKey(&pBindingInfo->m_oPublKey));

	/* Encrypt the content key with the blackbox public key */
	ChkDR(DRM_PK_Encrypt(g_BBContext.CryptoContext.rgbCryptoContext, &pBindingInfo->m_oPublKey, ClearContentKey, pBindingInfo->m_rgbContentKey));
	
	/* Calculate the content key's checksum */
	SetCheckSum(ClearContentKey + 1, keyLength, pBindingInfo->m_rgbChecksum);

	/* Generate the license server pub/private keys */
	ChkDR(DRM_PK_GenKeyPair(g_BBContext.CryptoContext.rgbCryptoContext, &pBindingInfo->m_oLsPublKey, &lsPrvKey));

	/* Sign the content key with the license server private key */
	if (keyLength) {
		ChkDR(DRM_PK_Sign(g_BBContext.CryptoContext.rgbCryptoContext, &lsPrvKey, ClearContentKey + 1, keyLength, pBindingInfo->m_rgbSignature));
	}

	/* We have a valid binding info blob now. */

	/* Intentionally modify the binding info for other tests... */
	if (!DX_VOS_StrCmp(argv[1], "BAD_CHECKSUM")) /* flip a byte in the checksum */
		pBindingInfo->m_rgbChecksum[1] ^= 0xff;

	else if (!DX_VOS_StrCmp(argv[1], "BAD_SIGNATURE")) /* flip a byte in the signatrue */
		pBindingInfo->m_rgbSignature[2] ^= 0xff;

	else if (!DX_VOS_StrCmp(argv[1], "BAD_CONTENTKEY")) /* flip a byte in the encrypted content key */
		pBindingInfo->m_rgbContentKey[1] ^= 0xff;

	else if (!DX_VOS_StrCmp(argv[1], "BAD_LSPUBKEY")) /* flip a byte in the license server pub key */
		pBindingInfo->m_oLsPublKey.y[2] ^= 0xff;

	else if (!DX_VOS_StrCmp(argv[1], "BAD_BBPUBKEY")) /* flip a byte in the blackbox pub key */
		pBindingInfo->m_oPublKey.y[2] ^= 0xff;

	else if (!DX_VOS_StrCmp(argv[1], "WRONG_LSPUBKEY")) { /* license server pubkey is different. Someone tries to re-sign the license */
		ChkDR(DRM_PK_GenKeyPair(g_BBContext.CryptoContext.rgbCryptoContext, &pBindingInfo->m_oLsPublKey, &lsPrvKey));
	}
	
ErrorExit:
	return dr;
}

/*
	Test API DRM_BBX_CanBind
	argv[0]: index to the global BindingInfo created by Test_CreateBindingInfo
	argv[1]: blackbox context: NULL, NORMAL
*/
DRM_RESULT Test_API_CanBind(long argc, char **argv)
{
	DRM_RESULT dr;
	DRM_BINDING_INFO *pBindingInfo;
	DRM_BB_CONTEXT *pBBContext;
	
	ChkArg(argc == 2);
	ChkDR(ArgToBindingInfo(argv[0], &pBindingInfo));
	pBBContext = argv[1]? &g_BBContext: NULL;

	dr = DRM_BBX_CanBind(pBindingInfo, 1, pBBContext)? DRM_SUCCESS: DRM_S_FALSE;
ErrorExit:
	return dr;
}
