#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <openssl/aes.h>
#include <string.h>
#include <math.h>
#include <errno.h>

/* read buffer size, must be reasonable value for online entropy test */
#define BUFFLEN 4096

/* check if we call the AES encryption function correctly, and if it works */
int aes_selftest()
{
	unsigned char *key = (unsigned char *)
		"\x2b\x7e\x15\x16\x28\xae\xd2\xa6\xab\xf7\x15\x88\x09\xcf\x4f\x3c";
	unsigned char *plain = (unsigned char *)
		"\x32\x43\xf6\xa8\x88\x5a\x30\x8d\x31\x31\x98\xa2\xe0\x37\x07\x34";
	unsigned char *ref_cipher = (unsigned char *)
		"\x39\x25\x84\x1d\x02\xdc\x09\xfb\xdc\x11\x85\x97\x19\x6a\x0b\x32";
	unsigned char cipher[16];
	AES_KEY ks;

	AES_set_encrypt_key(key, 128, &ks);
	memset(cipher, 0, 16);

	AES_encrypt(plain, cipher, &ks);

	if( !memcmp(cipher, ref_cipher, 16) )
		return 0;
	return -1;	
}

/* compute entropy per byte from relative frequencies */
double byte_ent(unsigned char *buff, int n)
{
	int counts[0x100], k, j;
	double z, ent;

	memset(counts, 0, sizeof(counts));
	for(k = 0; k < n; k++)
		counts[buff[k]]++;

	ent = 0.;
	for(j = 0; j < 0x100; j++)
	{
		if( counts[j] != 0 )
		{
			z = (double)counts[j];
			ent -= z * log(z);
		}	
	}
	ent /= (double)n;
	ent += log((double)n);
	ent /= log(2.);
	return ent;
}

/* read data into buffer */
int read_block(int fd, unsigned char *buff)
{
	int n, ret;

	n = 0;
	while( n != BUFFLEN )
	{
		ret = read(fd, buff + n, BUFFLEN - n);
		if( ret == 0 )
			break;
		else if( ret < 0 )
		{
			perror("read()");
			exit(errno);
		}
		n += ret;
	}
	return n;
}

int main(int argc, char **argv)
{
	AES_KEY ks;
	unsigned char ibuff[BUFFLEN], obuff[BUFFLEN];
	int k, nx, ji, jo;
	double ent, ent_min;
	
	k = 1;
	if( argc != 3
	 || sscanf(argv[k++], "%d", &nx) != 1
	 || sscanf(argv[k++], "%lf", &ent_min) != 1 )
	{
		fprintf(stderr, "Syntax:  %s <nx> <ent_min>\n", *argv);
		exit(EINVAL);
	}
	if( nx < 1 || nx >= 16 )
	{
		fprintf(stderr, "Error: <nx> must be in range 1 .. 16\n");
		exit(EINVAL);
	}
	if( ent_min < 0. || ent_min > 8. )
	{
		fprintf(stderr, "Error: <ent_min> must be in range 0.0 to 8.0\n");
		exit(EINVAL);
	}

	if( aes_selftest() != 0 )
	{
		fprintf(stderr, "*** AES self test failed!\n");
		exit(-1);
	}

	if( read_block(0, ibuff) != BUFFLEN )
		exit(errno);

	for(ji = 16; ji < BUFFLEN; ji++)
		ibuff[ji % 16] ^= ibuff[ji];
	AES_set_encrypt_key(ibuff, 128, &ks);

	for(;;)
	{
		if( read_block(0, ibuff) != BUFFLEN )
			break;
		
		ent = byte_ent(ibuff, BUFFLEN);	
		if( ent < ent_min )
		{
			fprintf(stderr, "ALARM: Low entropy, %.6lf < %.6lf\n", ent, ent_min);
			exit(EINVAL);
		}

		for(ji = 0, jo = 0; ji < BUFFLEN; ji += 16, jo += nx)
			AES_encrypt(ibuff + ji, obuff + jo, &ks);

		if( write(1, obuff, jo) != jo )
		{
			perror("write()");
			exit(errno);
		}
	}
	return 0;
}