/*
 *  Cryptographic "Data-as-Key" Feedback Analyser
 *
 *  Copyright (C) 2005 Luke Kenneth Casson Leighton
 *
 *  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.
 *
 *  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.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */
/*#define TEST_HISTOGRAM*/

#ifdef TEST_RIJNDAEL
#include "aes.h"
#endif

#ifdef TEST_SERPENT
#include "serpent.h"
#endif


#include <string.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "tomcrypt.h"

/*
 * the basic premise of this program is as follows:
 *
 * for a given initial key "k" and initial data "buf",
 * repeat num_tests times:
 *
 * 	o = E(k, buf); k = o; buf = xor(buf, o)
 *
 * if the algorithm is any good, then using the key as the
 * data for the next block will show no signs of non-randomness
 * in the output.  this just seems like good common sense.
 *
 * (and Rijndael/128 shows problems on blocks as quickly as 5-10 loops)
 *
 * signs of non-randomness aren't such a big hairy deal... as long as
 * there's no correlation between bits being set in the initial key
 * and bits being set in the initial data... and that's what the
 * histogram is there to find out.
 *
 * use gnuplot in weird and wonderful ways to look at the plots in 3d...
 *
 */

#define NUM_BLOCKS 2

extern int num_bits(unsigned char *data, int len);
extern double erfc(double x);

void bit_set(unsigned char *d, int bit)
{
	d[bit/8] |= 1<<(bit&0x7);
}

static void xor(unsigned char *key2, unsigned char *key1, unsigned char *xorval)
{
	int i;
	for (i = 0; i < 16; i++)
		key2[i] = key1[i] ^ xorval[i];
}

/* okay, this is really simple.
 * repeat num_tests times:
 *
 * 	o = E(k, buf); k = o; buf = xor(buf, o)
 *
 */
int analyse( unsigned char *d, unsigned char *k, unsigned char *xordata )
{
	int j;
	unsigned char buf[16];
	unsigned char out[16];
	unsigned char key[16];
	unsigned char test_data[16*NUM_BLOCKS]; /* make this 32..64... for larger tests */
	int count;
	double sobs;
	double testval;
	double pval;
	double sqrt2 = 1.41421356237309504880;
	double f;

	memcpy(key, k, 16);
	memcpy(buf, d, 16);

	for (j = 0; j < sizeof(test_data); j += 16)
	{
		/* o = E(k, buf); k = o; buf = xor(buf, o) */
#ifdef TEST_RIJNDAEL
		aes_context ctx;
		aes_set_key( &ctx, key, 128 );
		aes_encrypt( &ctx, buf, out );
#endif
#ifdef TEST_SERPENT
		unsigned int ctx[140];
		serpent_set_key( ctx, key, 128 );
		serpent_encrypt( &ctx, buf, out);
#endif
		memcpy(key, out, 16);
		xor(buf, out, buf);
		xor(buf, buf, xordata);

		/* output - concatenate data for testing... */
		memcpy(&test_data[j], out, 16);
	}

	/* count the number of bits, create a p-value, blah blah */
	count = num_bits(test_data, sizeof(test_data));
	testval = (count - ((double)sizeof(test_data)*4)) * 2;
	sobs = fabs(testval) / sqrtl(sizeof(test_data));
	f = sobs / sqrt2;
	pval = erfc(f);

	if (pval < 0.01)
		return 1;
	/*printf("testval: %10.f pval: %f\n", testval, pval);*/

	return 0;
}

struct datapoint
{
	unsigned char i;
	unsigned char j;
	int bitbucket;
	struct datapoint *next;
};

static int bitbucket[128][128];
static struct datapoint *sortedbucket = NULL;

void dump_bitbucket(void)
{
	fwrite(bitbucket, sizeof(bitbucket), 1, stdout);
}
			

void print_bitbucket(int num_samples)
{
	int i, j;
	for (i = 0; i < 128; i++)
	{
		int print = 0;
		for (j = 0; j < 128; j++)
		{
			if (bitbucket[i][j] != 0)
			{
				print = 1;
				break;
			}
		}
		if (!print)
			continue;

		for (j = 0; j < 128; j++)
		{
			if (bitbucket[i][j] != 0)
				printf("%3d ", (int)(100.*128.*128.*(double)bitbucket[i][j]/(double)num_samples));
		}
		printf("\n");
	}
}


int test(struct datapoint *i1, struct datapoint *i2)
{
	if (i1->bitbucket > i2->bitbucket)
		return 1;
	/*if (i1->bitbucket == i2->bitbucket && i1->j > i2->j)
		return 1;*/
	return 0;
}

void insert_datapoint(int i, int j, int bitbucket)
{
	struct datapoint *d = malloc(sizeof(struct datapoint));
	struct datapoint *start = sortedbucket;

	d->i = i;
	d->j = j;
	d->bitbucket = bitbucket;
	d->next = NULL;

	if (start == NULL)
	{
		sortedbucket = d;
		return;
	}

	if (test(sortedbucket, d))
	{
		d->next = sortedbucket;
		sortedbucket = d;
		return;
	}

	while (start->next != NULL)
	{
		if (test(start->next, d))
		{
			d->next = start->next;
			start->next = d;
			return;
		}
		start = start->next;
	}
	start->next = d;
}

static unsigned char rnumsj[128];
static unsigned char rnumsi[128];

void print_sortedbucket(int num_samples)
{
	int i, j;
	struct datapoint *d;
	while (sortedbucket != NULL)
	{
		struct datapoint *next = sortedbucket->next;
		free(sortedbucket);
		sortedbucket = next;
	}

	for (i = 0; i < 128; i++)
	{
		for (j = 0; j < 128; j++)
		{
			insert_datapoint(i, j, bitbucket[i][j]);
		}
	}
	d = sortedbucket;
	while (d != NULL)
	{
#if 0
		printf("i:%2x j:%2x "
#if 0
				"ri:%2x rj: %2x "
				"di:%3d dj:%3d "
#endif
				"bb: %5d\n",
				d->i, d->j,
#if 0
				rnumsi[d->i], rnumsj[d->j],
				rnumsj[d->j] - d->i,
				d->j - rnumsi[d->i],
#endif
				d->bitbucket);
#else
		printf("%d.\t%d.\t%d.\n", d->i, d->j, d->bitbucket);
#endif
		d = d->next;
	}
	fflush(stdout);
}

void swap(unsigned char *rnums, unsigned char i, unsigned char j)
{
	unsigned char tmp = rnums[i];
	rnums[i] = rnums[j];
	rnums[j] = tmp;
}

void init_rnums(unsigned char *rnums)
{
	unsigned char i;
	for (i = 0; i < 128; i++)
		rnums[i] = i;
}

void init_ij_rnums(void)
{
	init_rnums(rnumsj);
	init_rnums(rnumsi);
}

void print_ij_rnums(void)
{
	unsigned char i;
	for (i = 0; i < 128; i++)
		printf("%3d %3d\n", rnumsi[i], rnumsj[i]);
}

void diffcount( int *icount, int *jcount)
{
	struct datapoint *d;
	d = sortedbucket;
	*jcount = 0;
	*icount = 0;
	while (d != NULL)
	{
		int idiff = abs(d->j - rnumsi[d->i]);
		int jdiff = abs(d->i - rnumsj[d->j]);

		*icount += idiff;
		*jcount += jdiff;

		d = d->next;
	}
}

#if 0

static int prevdiffi = 1<<30;
static int prevdiffj = 1<<30;

void get_best_diffcount(void)
{
	int counti = 0;
	int countj = 0;

	int r = rand();

	int i_j = (r>>0) & 127;
	int i_i = (r>>7) & 127;
	int j_j = (r>>14) & 127;
	int j_i = (r>>21) & 127;

	static int rnums_inited;
	
	if (!rnums_inited)
	{
		init_ij_rnums();
		rnums_inited = 1;
	}

	swap(rnumsi, i_i, i_j);
	swap(rnumsj, j_i, j_j);

	diffcount(&counti, &countj);

	if (countj >= prevdiffj)
	{
		swap(rnumsj, j_i, j_j);
	}
	else
	{
		prevdiffj = countj;
	}
	if (counti >= prevdiffi)
	{
		swap(rnumsi, i_i, i_j);
	}
	else
	{
		prevdiffi = counti;
	}

}

void repeat_best_count(int goes)
{
	int i;
	for (i = 0; i < goes; i++)
		get_best_diffcount();
}
#endif

extern int gen_rand_impl(unsigned char *, unsigned int len);

static void pseud(unsigned char *d)
{
	gen_rand_impl(d, 16);
}

int main( void )
{
	int k, i, j;
	unsigned char buf[16];
	unsigned char key[16];
	unsigned char input[16];
	int total = 0;

	for (k = 0; k < 204800; k++) /* the more, the merrier */
	{
		/*pseud(input);  paranoia: get a decent PRNG, dude */
		memset( input, 0x0, 16 );
		/*bit_set(input, k);*/
		(*(int*)input) = k;

		for (i = 0; i < 128; i++)
		{
			memset( buf, 0x0, 16 );
			bit_set(buf, i);
			/*bit_set(buf, k % 128);*/

			/*printf("%d-> ", i);*/
			for (j = 0; j < 128; j++)
			{
				int num_fails;
				memset( key, 0x0, 16 );
				bit_set(key, j);

				/* get the number of times that the
				 * p-value (see NIST SP800-22b.pdf)
				 * fell below 0.01 */
				num_fails = analyse(input, key, buf);
				/*printf("i: %d j: %d num_fails: %d/%d\n",
						i, j, num_fails);*/

				/* create a histogram of the failures, using
				 * the key bit index and the data bit index */
				bitbucket[i][j] += num_fails;
				total += num_fails;
			}
		}
		if (k % 16 == 0)
		{
			fprintf(stderr, "%4d ", k);
			fflush(stderr);
		}
	}
	/*print_bitbucket(total);*/
	/*print_sortedbucket(total);*/
	/*dump_bitbucket();*/
	print_sortedbucket(total);
#if 0
	repeat_best_count(10000);
	print_sortedbucket(total);
	print_ij_rnums();
#endif
	return( 0 );
}

