/*  test_ctr_prng.c - TinyCrypt implementation of some CTR-PRNG tests */

/*
 * Copyright (c) 2016, Chris Morrison, All Rights Reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * * Redistributions of source code must retain the above copyright notice, this
 *   list of conditions and the following disclaimer.
 *
 * * Redistributions in binary form must reproduce the above copyright notice,
 *   this list of conditions and the following disclaimer in the documentation
 *   and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

/*
  DESCRIPTION
  This module tests the CTR-PRNG routines
*/

#include <tinycrypt/ctr_prng.h>
#include <tinycrypt/aes.h>
#include <tinycrypt/constants.h>
#include <test_utils.h>


#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/* utility function to convert hex character representation to their nibble (4 bit) values */
static uint8_t nibbleFromChar(char c)
{
	if(c >= '0' && c <= '9') return c - '0';
	if(c >= 'a' && c <= 'f') return c - 'a' + 10U;
	if(c >= 'A' && c <= 'F') return c - 'A' + 10U;
	return 255U;
}

/* 
 * Convert a string of characters representing a hex buffer into a series of 
 * bytes of that real value 
 */
uint8_t *hexStringToBytes(char *inhex)
{
	uint8_t *retval;
	uint8_t *p;
	int len, i;
	
	len = strlen(inhex) / 2;
	retval = (uint8_t *)malloc(len+1);
	for(i=0, p = (uint8_t *) inhex; i<len; i++) {
		retval[i] = (nibbleFromChar(*p) << 4) | nibbleFromChar(*(p+1));
		p += 2;
	}
	retval[len] = 0;
	return retval;
}

typedef struct {
	char * entropyString;
	char * personalizationString;  /* may be null */
	char * additionalInputString1; /* may be null */
	char * additionalInputString2; /* may be null */
	char * expectedString;
} PRNG_Vector;

/* vectors taken from NIST CAVS 14.3 CTR_DRBG.rsp */
PRNG_Vector vectors[] = {
	/*
	 * AES-128 no df, PredictionResistance = False, EntropyInputLen = 256,
	 * NonceLen = 0, PersonalizationStringLen = 0, AdditionalInputLen = 0,
	 * ReturnedBitsLen = 512
	 */
	{  /* Count 0 */
		"ce50f33da5d4c1d3d4004eb35244b7f2cd7f2e5076fbf6780a7ff634b249a5fc",
		0,
		0,
		0,
		"6545c0529d372443b392ceb3ae3a99a30f963eaf313280f1d1a1e87f9db373d361e75d18018266499cccd64d9bbb8de0185f213383080faddec46bae1f784e5a",
	},

	{ /* Count 1 */
		"a385f70a4d450321dfd18d8379ef8e7736fee5fbf0a0aea53b76696094e8aa93",
		0,
		0,
		0,
		"1a062553ab60457ed1f1c52f5aca5a3be564a27545358c112ed92c6eae2cb7597cfcc2e0a5dd81c5bfecc941da5e8152a9010d4845170734676c8c1b6b3073a5",
	},

	/*
	 * AES-128 no df, PredictionResistance = False, EntropyInputLen = 256,
	 * NonceLen = 0, PersonalizationStringLen = 0, AdditionalInputLen = 256,
	 * ReturnedBitsLen = 512
	 */
	{ /* Count 0 */
		"6bd4f2ae649fc99350951ff0c5d460c1a9214154e7384975ee54b34b7cae0704",
		0,
		"ecd4893b979ac92db1894ae3724518a2f78cf2dbe2f6bbc6fda596df87c7a4ae",
		"b23e9188687c88768b26738862c4791fa52f92502e1f94bf66af017c4228a0dc",
		"5b2bf7a5c60d8ab6591110cbd61cd387b02de19784f496d1a109123d8b3562a5de2dd6d5d1aef957a6c4f371cecd93c15799d82e34d6a0dba7e915a27d8e65f3",
	},

	{ /* Count 1 */
		"e2addbde2a76e769fc7aa3f45b31402f482b73bbe7067ad6254621f06d3ef68b",
		0,
		"ad11643b019e31245e4ea41f18f7680458310580fa6efad275c5833e7f800dae",
		"b5d849616b3123c9725d188cd0005003220768d1200f9e7cc29ef6d88afb7b9a",
		"132d0d50c8477a400bb8935be5928f916a85da9ffcf1a8f6e9f9a14cca861036cda14cf66d8953dab456b632cf687cd539b4b807926561d0b3562b9d3334fb61",
	},

	/*
	 * AES-128 no df, PredictionResistance = False, EntropyInputLen = 256,
	 * NonceLen = 0, PersonalizationStringLen = 256, AdditionalInputLen = 0,
	 * ReturnedBitsLen = 512
	 */
	{  /* Count 0 */
		"cee23de86a69c7ef57f6e1e12bd16e35e51624226fa19597bf93ec476a44b0f2",
		"a2ef16f226ea324f23abd59d5e3c660561c25e73638fe21c87566e86a9e04c3e",
		0,
		0,
		"2a76d71b329f449c98dc08fff1d205a2fbd9e4ade120c7611c225c984eac8531288dd3049f3dc3bb3671501ab8fbf9ad49c86cce307653bd8caf29cb0cf07764",
	},

	{ /* Count 1 */
		"b09eb4a82a39066ec945bb7c6aef6a0682a62c3e674bd900297d4271a5f25b49",
		"a3b768adcfe76d61c972d900da8dffeeb2a42e740247aa719ed1c924d2d10bd4",
		0,
		0,
		"5a1c26803f3ffd4daf32042fdcc32c3812bb5ef13bc208cef82ea047d2890a6f5dcecf32bcc32a2585775ac5e1ffaa8de00664c54fe00a7674b985619e953c3a",
	},

	/*
	 * AES-128 no df, PredictionResistance = False, EntropyInputLen = 256,
	 * NonceLen = 0, PersonalizationStringLen = 256, AdditionalInputLen = 256,
	 * ReturnedBitsLen = 512
	 */
	{  /* Count 0 */
		"50b96542a1f2b8b05074051fe8fb0e45adbbd5560e3594e12d485fe1bfcb741f",
		"820c3030f97b3ead81a93b88b871937278fd3d711d2085d9280cba394673b17e",
		"1f1632058806d6d8e231288f3b15a3c324e90ccef4891bd595f09c3e80e27469",
		"5cadc8bfd86d2a5d44f921f64c7d153001b9bdd7caa6618639b948ebfad5cb8a",
		"02b76a66f103e98d450e25e09c35337747d987471d2b3d81e03be24c7e985417a32acd72bc0a6eddd9871410dacb921c659249b4e2b368c4ac8580fb5db559bc",
	},

	{ /* Count 1 */
		"ff5f4b754e8b364f6df0c5effba5f1c036de49c4b38cd8d230ee1f14d7234ef5",
		"994eb339f64034005d2e18352899e77df446e285c3430631d557498aac4f4280",
		"e1824832d5fc2a6dea544cac2ab73306d6566bde98cc8f9425d064b860a9b218",
		"c08b42433a78fd393a34ffc24724d479af08c36882799c134165d98b2866dc0a",
		"1efa34aed07dd57bde9741b8d1907d28e8c1ac71601df37ef4295e6ffb67f6a1c4c13e5def65d505e2408aeb82948999ca1f9c9113b99a6b59ff7f0cc3dc6e92",
	},

	/*
	 * AES-128 no df, PredictionResistance = False, EntropyInputLen = 256,
	 * NonceLen = 0, PersonalizationStringLen = 0, AdditionalInputLen = 0,
	 * ReturnedBitsLen = 512
	 */
	{  /* Count 0 */
		"69a09f6bf5dda15cd4af29e14cf5e0cddd7d07ac39bba587f8bc331104f9c448",
		0,
		0,
		0,
		"f78a4919a6ec899f7b6c69381febbbe083315f3d289e70346db0e4ec4360473ae0b3d916e9b6b964309f753ed66ae59de48da316cc1944bc8dfd0e2575d0ff6d",
	},

	{ /* Count 1 */
		"80bfbd340d79888f34f043ed6807a9f28b72b6644d9d9e9d777109482b80788a",
		0,
		0,
		0,
		"80db048d2f130d864b19bfc547c92503e580cb1a8e1f74f3d97fdda6501fb1aa81fcedac0dd18b6ccfdc183ca28a44fc9f3a08834ba8751a2f4495367c54a185",
	},

	/*
	 * AES-128 no df, PredictionResistance = False, EntropyInputLen = 256,
	 * NonceLen = 0, PersonalizationStringLen = 0, AdditionalInputLen = 256,
	 * ReturnedBitsLen = 512
	 */
	{ /* Count 0 */
		"7f40804693552e317523fda6935a5bc814353b1fbb7d334964ac4d1d12ddccce",
		0,
		"95c04259f64fcd1fe00c183aa3fb76b8a73b4d1243b800d770e38515bc41143c",
		"5523102dbd7fe1228436b91a765b165ae6405eb0236e237afad4759cf0888941",
		"1abf6bccb4c2d64e5187b1e2e34e493eca204ee4eef0d964267e38228f5f20efba376430a266f3832916d0a45b2703f46401dfd145e447a0a1667ebd8b6ee748",
	},

	{ /* Count 1 */
		"350df677409a1dc297d01d3716a2abdfa6272cd030ab75f76839648582b47113",
		0,
		"ba5709a12ae6634a5436b7ea06838b48f7b847a237f6654a0e27c776ebee9511",
		"f1b2c717c5e3a934127e10471d67accc65f4a45010ca53b35f54c88833dbd8e7",
		"1ef1ea279812e8abe54f7ffd12d04c80ae40741f4ccfe232a5fba3a78dfd3e2ed419b88ee9188df724160cbb3aea0f276e84a3c0ff01e3b89fe30ebcfa64cb86",
	},

	/*
	 * AES-128 no df, PredictionResistance = False, EntropyInputLen = 256,
	 * NonceLen = 0, PersonalizationStringLen = 256, AdditionalInputLen = 0,
	 * ReturnedBitsLen = 512
	 */
	{  /* Count 0 */
		"3fef762f0aa0677f61c65d749eeb10b013ff68ccc6314f150cfee752dcd8f987",
		"f56db099240c7590dac396372b8737404d418b2864a3df96a8a397967245735f",
		0,
		0,
		"af0afe0837442136fbb1959a1c91a9291c1d8188ede07c67d0e4dd6541303415e7a67999c302ba0df555324c26077514592a9b6db6be2f153fad2250161164e4",
	},

	{ /* Count 1 */
		"3eebe77db4631862e3eb7e39370515b8baa1cdd71a5b1b0cda79c14d0b5f48ea",
		"4be56a9b9c21242739c985ef12aa4d98e8c7da07c4c1dc6829f2e06833cfa148",
		0,
		0,
		"be9e18a753df261927473c8bb5fb7c3ea6e821df5ab49adc566a4ebf44f75fa825b1f9d8c154bcd469134c0bb688e07e3c3e45407ca350d540e1528cc2e64068",
	},

	/*
	 * AES-128 no df, PredictionResistance = False, EntropyInputLen = 256,
	 * NonceLen = 0, PersonalizationStringLen = 256, AdditionalInputLen = 256,
	 * ReturnedBitsLen = 512
	 */
	{  /* Count 0 */
		"c129c2732003bbf1d1dec244a933cd04cb47199bbce98fe080a1be880afb2155",
		"64e2b9ac5c20642e3e3ee454b7463861a7e93e0dd1bbf8c4a0c28a6cb3d811ba",
		"f94f0975760d52f47bd490d1623a9907e4df701f601cf2d573aba803a29d2b51",
		"6f99720b186e2028a5fcc586b3ea518458e437ff449c7c5a318e6d13f75b5db7",
		"7b8b3378b9031ab3101cec8af5b8ba5a9ca2a9af41432cd5f2e5e19716140bb219ed7f4ba88fc37b2d7e146037d2cac1128ffe14131c8691e581067a29cacf80",
	},

	{ /* Count 1 */
		"7667643670254b3530e80a17b16b22406e84efa6a4b5ceef3ebc877495fc6048",
		"40b92969953acde756747005117e46eff6893d7132a8311ffb1062280367326b",
		"797a02ffbe8ff2c94ed0e5d39ebdc7847adaa762a88238242ed8f71f5635b194",
		"d617f0f0e609e90d814192ba2e5214293d485402cdf9f789cc78b05e8c374f18",
		"e8d6f89dca9825aed8927b43187492a98ca8648db30f0ac709556d401a8ac2b959c81350fc64332c4c0deb559a286a72e65dbb462bd872f9b28c0728f353dc10",
	}
};

static unsigned int executePRNG_TestVector(PRNG_Vector vector, unsigned int idx)
{
	unsigned int result = TC_PASS;
	uint8_t * entropy    = hexStringToBytes(vector.entropyString);
	unsigned int  entropylen = strlen(vector.entropyString) / 2U;

	uint8_t * expected    = hexStringToBytes(vector.expectedString);
	unsigned int  expectedlen = strlen(vector.expectedString) / 2U;

	uint8_t * personalization   = 0;
	unsigned int  plen              = 0U;

	uint8_t * additional_input1 = 0;
	unsigned int  additionallen1    = 0U;

	uint8_t * additional_input2 = 0;
	unsigned int  additionallen2    = 0U;

	uint8_t * output = (uint8_t *)malloc(expectedlen);

	unsigned int i;
	TCCtrPrng_t ctx;

	if (0 != vector.personalizationString) {
		personalization = hexStringToBytes(vector.personalizationString);
		plen = strlen(vector.personalizationString) / 2U;
	}

	if (0 != vector.additionalInputString1) {
		additional_input1 = hexStringToBytes(vector.additionalInputString1);
		additionallen1 = strlen(vector.additionalInputString1) / 2U;
	}

	if (0 != vector.additionalInputString2) {
		additional_input2 = hexStringToBytes(vector.additionalInputString2);
		additionallen2 = strlen(vector.additionalInputString2) / 2U;
	}

	(void)tc_ctr_prng_init(&ctx, entropy, entropylen, personalization, plen);

	(void)tc_ctr_prng_generate(&ctx, additional_input1, additionallen1, output, expectedlen);
	(void)tc_ctr_prng_generate(&ctx, additional_input2, additionallen2, output, expectedlen);

	for (i = 0U; i < expectedlen; i++) {
		if (output[i] != expected[i]) {
			TC_ERROR("CTR PRNG test #%d failed\n", idx);
			result = TC_FAIL;
			break;
		}
	}

	free(entropy);
	free(expected);
	free(personalization);
	free(additional_input1);
	free(additional_input2);
	free(output);

	return result;
}

static int test_reseed(void)
{
	int result = TC_PASS;
	uint8_t entropy[32U] = {0U}; /* value not important */
	uint8_t additional_input[32] = {0U};
	uint8_t output[32];
	TCCtrPrng_t ctx;
	int ret;
	unsigned int i;

	(void)tc_ctr_prng_init(&ctx, entropy, sizeof entropy, 0, 0U);

	/* force internal state to max allowed count */
	ctx.reseedCount = 0x1000000000000ULL;

	ret = tc_ctr_prng_generate(&ctx, 0, 0, output, sizeof output);
	if (1 != ret) {
		result = TC_FAIL;
		goto exitTest;
	}

	/* expect further attempts to fail due to reaching reseed threshold */
	ret = tc_ctr_prng_generate(&ctx, 0, 0, output, sizeof output);
	if (-1 != ret) {
		result = TC_FAIL;
		goto exitTest;
	}

	/* reseed and confirm generate works again */
	/* make entropy different from original value - not really important for the purpose of this test */
	memset(entropy, 0xFF, sizeof entropy);
	ret = tc_ctr_prng_reseed(&ctx, entropy, sizeof entropy, additional_input, sizeof additional_input);
	if (1 != ret) {
		result = TC_FAIL;
		goto exitTest;
	}

	ret = tc_ctr_prng_generate(&ctx, 0, 0, output, sizeof output);
	if (1 != ret) {
		result = TC_FAIL;
		goto exitTest;
	}

	/* confirm entropy and additional_input are being used correctly */
	/* first, entropy only */
	memset(&ctx, 0x0, sizeof ctx);
	for (i = 0U; i < sizeof entropy; i++) {
		entropy[i] = i;
	}
	ret = tc_ctr_prng_reseed(&ctx, entropy, sizeof entropy, 0, 0U);
	if (1 != ret) {
		result = TC_FAIL;
		goto exitTest;
	}
	{
		uint8_t expectedV[] =
			{0x7EU, 0xE3U, 0xA0U, 0xCBU, 0x6DU, 0x5CU, 0x4BU, 0xC2U,
			 0x4BU, 0x7EU, 0x3CU, 0x48U, 0x88U, 0xC3U, 0x69U, 0x70U};
		for (i = 0U; i < sizeof expectedV; i++) {
			if (ctx.V[i] != expectedV[i]) {
				result = TC_FAIL;
				goto exitTest;
			}
		}
	}

	/* now, entropy and additional_input */
	memset(&ctx, 0x0, sizeof ctx);
	for (i = 0U; i < sizeof additional_input; i++) {
		additional_input[i] = i * 2U;
	}
	ret = tc_ctr_prng_reseed(&ctx, entropy, sizeof entropy, additional_input, sizeof additional_input);
	if (1 != ret) {
		result = TC_FAIL;
		goto exitTest;
	}
	{
		uint8_t expectedV[] =
			{0x5EU, 0xC1U, 0x84U, 0xEDU, 0x45U, 0x76U, 0x67U, 0xECU,
			 0x7BU, 0x4CU, 0x08U, 0x7EU, 0xB0U, 0xF9U, 0x55U, 0x4EU};
		for (i = 0U; i < sizeof expectedV; i++) {
			if (ctx.V[i] != expectedV[i]) {
				result = TC_FAIL;
				goto exitTest;
			}
		}
	}

	exitTest:
	if (TC_FAIL == result) {
		TC_ERROR("CTR PRNG reseed tests failed\n");
	}
	return result;
}

static int test_uninstantiate(void)
{
	unsigned int i;
	int result = TC_PASS;
	uint8_t entropy[32U] = {0U}; /* value not important */
	TCCtrPrng_t ctx;

	(void)tc_ctr_prng_init(&ctx, entropy, sizeof entropy, 0, 0U);

	tc_ctr_prng_uninstantiate(&ctx);
	/* show that state has been zeroised */
	for (i = 0U; i < sizeof ctx.V; i++) {
		if (0U != ctx.V[i]) {
			TC_ERROR("CTR PRNG uninstantiate tests failed\n");
			result = TC_FAIL;
			break;
		}
	}

	for (i = 0U; i < sizeof ctx.key.words / sizeof ctx.key.words[0]; i++) {
		if (0U != ctx.key.words[i]) {
			TC_ERROR("CTR PRNG uninstantiate tests failed\n");
			result = TC_FAIL;
			break;
		}
	}

	if (0U != ctx.reseedCount) {
		TC_ERROR("CTR PRNG uninstantiate tests failed\n");
		result = TC_FAIL;
	}

	return result;
}

static int test_robustness(void)
{
	int result = TC_PASS;
	int ret;
	uint8_t entropy[32U] = {0U}; /* value not important */
	uint8_t output[32];
	TCCtrPrng_t ctx;


	/* show that the CTR PRNG is robust to invalid inputs */
	tc_ctr_prng_uninstantiate(0);

	ret = tc_ctr_prng_generate(&ctx, 0, 0, 0, 0);
	if (0 != ret) {
		result = TC_FAIL;
		goto exitTest;
	}

	ret = tc_ctr_prng_generate(0, 0, 0, output, sizeof output);
	if (0 != ret) {
		result = TC_FAIL;
		goto exitTest;
	}

	ret = tc_ctr_prng_generate(0, 0, 0, 0, 0);
	if (0 != ret) {
		result = TC_FAIL;
		goto exitTest;
	}

	ret = tc_ctr_prng_reseed(&ctx, 0, 0, 0, 0);
	if (0 != ret) {
		result = TC_FAIL;
		goto exitTest;
	}

	/* too little entropy */
	ret = tc_ctr_prng_reseed(&ctx, entropy, (sizeof entropy) - 1U, 0, 0);
	if (0 != ret) {
		result = TC_FAIL;
		goto exitTest;
	}

	ret = tc_ctr_prng_reseed(0, entropy, sizeof entropy, 0, 0);
	if (0 != ret) {
		result = TC_FAIL;
		goto exitTest;
	}

	ret = tc_ctr_prng_reseed(0, 0, 0, 0, 0);
	if (0 != ret) {
		result = TC_FAIL;
		goto exitTest;
	}

	ret = tc_ctr_prng_init(&ctx, 0, 0, 0, 0);
	if (0 != ret) {
		result = TC_FAIL;
		goto exitTest;
	}

	/* too little entropy */
	ret = tc_ctr_prng_init(&ctx, entropy, (sizeof entropy) - 1U, 0, 0);
	if (0 != ret) {
		result = TC_FAIL;
		goto exitTest;
	}

	ret = tc_ctr_prng_init(0, entropy, sizeof entropy, 0, 0);
	if (0 != ret) {
		result = TC_FAIL;
		goto exitTest;
	}

	ret = tc_ctr_prng_init(0, 0, 0, 0, 0);
	if (0 != ret) {
		result = TC_FAIL;
		goto exitTest;
	}

	exitTest:
	if (TC_FAIL == result) {
		TC_ERROR("CTR PRNG reseed tests failed\n");
	}


	return result;
}

/*
 * Main task to test CTR PRNG
 */
int main(void)
{
	int result = TC_PASS;
	unsigned int i;
	TC_START("Performing CTR-PRNG tests:");
	for (i = 0U; i < sizeof vectors / sizeof vectors[0]; i++) {
		result = executePRNG_TestVector(vectors[i], i);
		if (TC_PASS != result) {
			goto exitTest;
		}
	}

	if (TC_PASS != test_reseed()) {
		goto exitTest;
	}

	if (TC_PASS != test_uninstantiate()) {
		goto exitTest;
	}

	if (TC_PASS != test_robustness()) {
		goto exitTest;
	}

	TC_PRINT("All CTR PRNG tests succeeded!\n");

	exitTest:
    	TC_END_RESULT(result);
    	TC_END_REPORT(result);

}