/*
 * Copyright (C) 2019 Tobias Brunner
 * HSR Hochschule fuer Technik Rapperswil
 *
 * This program is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the
 * Free Software Foundation; either version 2 of the License, or (at your
 * option) any later version.  See <http://www.fsf.org/copyleft/gpl.txt>.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
 * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * for more details.
 */

#include "test_suite.h"

#include <crypto/prf_plus.h>

static struct {
	chunk_t key;
	chunk_t seed;
	chunk_t iterations[10];
} counter_data[] = {
	{ .key = chunk_from_chars(0x0b,0x0b,0x0b,0x0b,0x0b,0x0b,0x0b,0x0b,
							  0x0b,0x0b,0x0b,0x0b,0x0b,0x0b,0x0b,0x0b,
							  0x0b,0x0b,0x0b,0x0b),
	  .seed = chunk_from_chars(0x48,0x69,0x20,0x54,0x68,0x65,0x72,0x65),
	  .iterations = {
		chunk_from_chars(0xb9,0xbd,0xc0),
		chunk_from_chars(0x89,0x88,0xb4,0xc2,0xb7,0x5a),
		chunk_from_chars(0xa9,0x3e,0x59,0x6a,0xc8,0x42,0x05),
		chunk_from_chars(0xfa,0x2d,0xdd,0xe1,0xbf,0x7a,0x25,0x72,
						 0x06,0x7b,0x00,0xe1,0x4b,0x23,0x77,0x32),
		chunk_from_chars(0x83,0x05,0x09,0x98,0x1a,0xd2,0xf9,0x4a),
		chunk_from_chars(0x8c,0x32,0xa4,0x7d,0xaa,0x22,0x55,0xb6),
		chunk_from_chars(0x60,0xc4,0x36,0x34,0x7a,0xe7,0x56,0xa6,
						 0xed,0xc0,0x23,0x47,0x7d,0x80),
		chunk_from_chars(0x95,0x90,0xe6,0x82,0xf6,0x1d,0x9c,0x04,
						 0xb0,0x6b,0x4a,0xd9,0x71,0xa3,0x4c,0x81,
						 0x47,0xfa,0x66,0x79),
		chunk_from_chars(0x2f,0xf1,0x43,0x4b,0x93,0xc7,0x22,0xb3,
						 0x2e,0x12,0xf4,0x88,0x32,0xeb,0xc1,0x5c,
						 0xe2,0x36,0x9c,0xe7,0x1f,0xe9,0xb7,0xb8,
						 0x1e,0x57,0x04,0xc1,0x4d,0x0f,0x52,0x80,
						 0xa6,0xec,0x62,0x6e,0x99,0x2d,0x7a,0x9f),
	  },
	},
};

START_TEST(test_vectors_counter)
{
	prf_plus_t *prf_plus;
	prf_t *prf;
	chunk_t *iter = counter_data[_i].iterations, out;

	prf = lib->crypto->create_prf(lib->crypto, PRF_HMAC_SHA2_256);
	ck_assert(prf->set_key(prf, counter_data[_i].key));
	prf_plus = prf_plus_create(prf, TRUE, counter_data[_i].seed);
	while (iter->ptr)
	{
		ck_assert(prf_plus->allocate_bytes(prf_plus, iter->len, &out));
		ck_assert_chunk_eq(*iter, out);
		chunk_free(&out);
		iter++;
	}
	prf_plus->destroy(prf_plus);
	prf->destroy(prf);
}
END_TEST

START_TEST(test_wrap)
{
	prf_plus_t *prf_plus;
	prf_t *prf;
	u_char buf[32];
	int i;

	prf = lib->crypto->create_prf(lib->crypto, PRF_HMAC_SHA2_256);
	ck_assert(prf->set_key(prf, counter_data[0].key));
	prf_plus = prf_plus_create(prf, TRUE, counter_data[0].seed);
	for (i = 1; i < 256; i++)
	{
		ck_assert(prf_plus->get_bytes(prf_plus, sizeof(buf), buf));
	}
	ck_assert(!prf_plus->get_bytes(prf_plus, sizeof(buf), buf));
	prf_plus->destroy(prf_plus);
	prf->destroy(prf);
}
END_TEST

static struct {
	chunk_t key;
	chunk_t seed;
	chunk_t iterations[10];
} classic_data[] = {
	{ .key = chunk_from_chars(0x0b,0x0b,0x0b,0x0b,0x0b,0x0b,0x0b,0x0b,
							  0x0b,0x0b,0x0b,0x0b,0x0b,0x0b,0x0b,0x0b,
							  0x0b,0x0b,0x0b,0x0b),
	  .seed = chunk_from_chars(0x48,0x69,0x20,0x54,0x68,0x65,0x72,0x65),
	  .iterations = {
		chunk_from_chars(0xb0,0x34,0x4c),
		chunk_from_chars(0x61,0xd8,0xdb,0x38,0x53,0x5c),
		chunk_from_chars(0xa8,0xaf,0xce,0xaf,0x0b,0xf1,0x2b),
		chunk_from_chars(0x88,0x1d,0xc2,0x00,0xc9,0x83,0x3d,0xa7,
						 0x26,0xe9,0x37,0x6c,0x2e,0x32,0xcf,0xf7),
		chunk_from_chars(0xd0,0x9a,0xe2,0x4b,0x3a,0x83,0xff,0xd4),
		chunk_from_chars(0xb1,0xef,0xa5,0x94,0x5c,0xc5,0xed,0x85),
		chunk_from_chars(0xb0,0xb2,0xcc,0x56,0xfc,0xf7,0x5d,0x23,
						 0xa0,0xa3,0x4c,0xa4,0xdb,0xff,),
		chunk_from_chars(0xea,0xfd,0xaa,0x6a,0x3b,0xf4,0x11,0x34,
						 0x24,0xe4,0x50,0x2d,0xf9,0x7a,0x76,0x93,
						 0x24,0xf6,0x11,0x24),
		chunk_from_chars(0x24,0x3b,0x99,0x6e,0x7d,0x0f,0x35,0x99,
						 0x88,0x79,0x73,0x6b,0xdb,0x70,0x65,0x9a,
						 0x6e,0xfa,0xd2,0x39,0x94,0x10,0xe6,0xce,
						 0x80,0x45,0x6e,0xb6,0x07,0x07,0x8f,0xe1,
						 0xc4,0x7c,0x6b,0x5e,0x81,0x65,0x47,0x8a),
	  },
	},
};

START_TEST(test_vectors_classic)
{
	prf_plus_t *prf_plus;
	prf_t *prf;
	chunk_t *iter = classic_data[_i].iterations, out;

	prf = lib->crypto->create_prf(lib->crypto, PRF_HMAC_SHA2_256);
	ck_assert(prf->set_key(prf, classic_data[_i].key));
	prf_plus = prf_plus_create(prf, FALSE, classic_data[_i].seed);
	while (iter->ptr)
	{
		ck_assert(prf_plus->allocate_bytes(prf_plus, iter->len, &out));
		ck_assert_chunk_eq(*iter, out);
		chunk_free(&out);
		iter++;
	}
	prf_plus->destroy(prf_plus);
	prf->destroy(prf);
}
END_TEST

Suite *prf_plus_suite_create()
{
	Suite *s;
	TCase *tc;

	s = suite_create("prf_plus");

	tc = tcase_create("counter");
	tcase_add_loop_test(tc, test_vectors_counter, 0, countof(counter_data));
	tcase_add_test(tc, test_wrap);
	suite_add_tcase(s, tc);

	tc = tcase_create("no counter");
	tcase_add_loop_test(tc, test_vectors_classic, 0, countof(classic_data));
	suite_add_tcase(s, tc);

	return s;
}