/**@@@+++@@@@******************************************************************
**
** Microsoft Windows Media
** Copyright (C) Microsoft Corporation. All rights reserved.
**
***@@@---@@@@******************************************************************
*/
#include "testdrmmanager.h"
#include "drmopcodes.h"

_XMBContext *g_pbDLRB = NULL;
DRM_DWORD g_cbDLRB = 10000;
DRM_STRING g_dstrDLRB;

const DRM_WCHAR wDLRB[]  = { ONE_WCHAR('d','\0'),  ONE_WCHAR('l','\0'),  ONE_WCHAR('r','\0'),  ONE_WCHAR('b','\0'), ONE_WCHAR('\0','\0') } ;

static DRM_RESULT _GetNextKID(DRM_CONST_STRING *pdstrKID)
{
	DRM_RESULT dr;
	DRM_KID kid;
	DRM_DWORD *pdw;
	
	ChkDR(DRM_UTL_DecodeKID(pdstrKID, &kid));
	pdw = (DRM_DWORD*)&kid;
	(*pdw)++;
	ChkDR(DRM_B64_EncodeW((DRM_BYTE*)&kid, sizeof(kid), (DRM_WCHAR*)pdstrKID->pwszString, &pdstrKID->cchString, 0));
ErrorExit:
	return dr;
}

static DRM_RESULT _AddKID(_XMBContext *pbXML, char *szStartingKID, int cKID)
{
	DRM_RESULT dr;
	int i;
	DRM_CONST_STRING dstrKID;

	MakeDRMString(&dstrKID, szStartingKID);

	for (i = 0; i < cKID; i++) {
		ChkDR(DRM_XMB_OpenNode(pbXML, &g_dstrTagKID));
		ChkDR(DRM_XMB_AddAttribute(pbXML, &g_dstrLabelValue, &dstrKID));
		ChkDR(DRM_XMB_CloseCurrNode(pbXML, NULL));

		ChkDR(_GetNextKID(&dstrKID));
	}
ErrorExit:
	DX_VOS_MemFree((DRM_WCHAR*)dstrKID.pwszString);
	return dr;
}


/* Generate a DLRB
	argv[0]: LGPUBKEY
	argv[1]: LG signing private key
	argv[2]: starting KID string
	argv[3]: number of KIDs in the DLRB
	argv[4]: optional flags to DLRB generation
*/
DRM_RESULT TestManagerMakeDeviceLicenseRevocationBlob(long argc, char **argv)
{
	DRM_RESULT dr;
       const DRM_CONST_STRING dstrDLRB = CREATE_DRM_STRING(wDLRB); 
	DRM_CONST_STRING dstrPubKey = {0}, dstrSignature;
	DRM_STRING dstrDataBlock;
	DRM_CRYPTO_CONTEXT crypto;
	DRM_SUBSTRING subStr;
	PRIVKEY prvKey;
	DRM_DWORD cb;
	DRM_WCHAR wszB64Sign[PK_ENC_SIGNATURE_B64LEN];

	ChkArg(argc >= 4 && argv[3]);

	if (!g_pbDLRB) {
		ChkMem(g_pbDLRB = (_XMBContext *)DX_VOS_MemMalloc(g_cbDLRB));
	}

	ChkDR(DRM_XMB_CreateDocument(g_cbDLRB, g_pbDLRB, &g_dstrTagDeviceLicenseRevocationRoot));
	ChkDR(DRM_XMB_AddAttribute(g_pbDLRB, &g_dstrAttributeType, &dstrDLRB));

	/* Add Data Node */
	ChkDR(DRM_XMB_OpenNode(g_pbDLRB, &g_dstrTagData));

		/* Add LGPUBKEY */
		ChkDR(DRM_XMB_OpenNode(g_pbDLRB, &g_dstrTagRevocationPubKey));
		MakeDRMString(&dstrPubKey, argv[0]);
		ChkDR(DRM_XMB_AddAttribute(g_pbDLRB, &g_dstrLabelValue, &dstrPubKey));
		ChkDR(DRM_XMB_CloseCurrNode(g_pbDLRB, &dstrDataBlock));

		ChkDR(_AddKID(g_pbDLRB, argv[2], OEM_atol(argv[3])));

	/* Get the signing private key */
	subStr.m_ich = 0;
	subStr.m_cch = DX_VOS_StrLen(argv[1]);
	cb = sizeof(PRIVKEY);
	ChkDR(DRM_B64_DecodeA(argv[1], &subStr, &cb, (DRM_BYTE*)&prvKey, 0));

	if (argc > 4 && !DX_VOS_StrCmp(argv[4], "BAD_SIGNATURE"))
		prvKey.x[0]++;

	ChkDR(DRM_XMB_SignAndCloseCurrNode(g_pbDLRB, &crypto, &prvKey, TRUE, NULL, wszB64Sign, NULL));

	/* Add SIGNATURE */
	ChkDR(DRM_XMB_OpenNode(g_pbDLRB, &g_dstrTagSignature));
		ChkDR(DRM_XMB_OpenNode(g_pbDLRB, &g_dstrTagHashAlg));
		ChkDR(DRM_XMB_AddAttribute(g_pbDLRB, &g_dstrAttributeType, &g_dstrSHA));
		ChkDR(DRM_XMB_CloseCurrNode(g_pbDLRB, NULL));

		ChkDR(DRM_XMB_OpenNode(g_pbDLRB, &g_dstrTagSignAlg));
		ChkDR(DRM_XMB_AddAttribute(g_pbDLRB, &g_dstrAttributeType, &g_dstrMSDRM));
		ChkDR(DRM_XMB_CloseCurrNode(g_pbDLRB, NULL));

		DSTR_FROM_PB(&dstrSignature, wszB64Sign, sizeof(wszB64Sign));
		ChkDR(DRM_XMB_OpenNode(g_pbDLRB, &g_dstrTagValue));
		ChkDR(DRM_XMB_AddData(g_pbDLRB, &dstrSignature));
		ChkDR(DRM_XMB_CloseCurrNode(g_pbDLRB, NULL));
	ChkDR(DRM_XMB_CloseCurrNode(g_pbDLRB, NULL));

	ChkDR(DRM_XMB_CloseDocument(g_pbDLRB, &g_dstrDLRB));
ErrorExit:
	DX_VOS_MemFree((DRM_WCHAR*)dstrPubKey.pwszString);
	return dr;
}

/*	Test API DRM_MGR_ProcessCommand(OpProcessLicenseRevocationList)
	argv[0]: index to the starting KID
	argv[1]: number of KIDs to revoke
	argv[2]: status of DLRB: NORMAL or NULL
	argv[3]: expected number of KIDs processed
	argv[4]: expected number of licenses deleted
*/
DRM_RESULT TestManagerProcessDeviceLicenseRevocation(long argc, char **argv)
{
	DRM_RESULT dr;
	DRM_DWORD dwcKidsProcessed, dwcLicsDeleted;

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

	ChkDR(DRM_MGR_ProcessCommand(
		argv[2]?g_pManagerContext_MNGR: NULL,
		OpProcessLicenseRevocationList,
		OEM_atol(argv[0]), OEM_atol(argv[1]), 0, 0,
		(DRM_BYTE*)g_dstrDLRB.pwszString, g_dstrDLRB.cchString * sizeof(DRM_WCHAR),
		(argv[3] && DX_VOS_StrCmp(argv[3], "IGNORE"))? &dwcKidsProcessed: NULL,
		(argv[4] && DX_VOS_StrCmp(argv[4], "IGNORE"))? &dwcLicsDeleted: NULL,
		NULL, NULL));

	/* Check return values */
	if (argv[3] && OEM_atol(argv[3]) != dwcKidsProcessed) {
		Log("Trace", "Number of processed KIDs does not match. Returning = %d. Expecting %d.", dwcKidsProcessed, OEM_atol(argv[3]));
		ChkDR(DRM_E_FAIL);
	}

	if (argv[4] && OEM_atol(argv[4]) != dwcLicsDeleted) {
		Log("Trace", "Number of deleted licenses does not match. Returning = %d. Expecting %d.", dwcLicsDeleted, OEM_atol(argv[4]));
		ChkDR(DRM_E_FAIL);
	}
ErrorExit:
	return dr;
}
