#include "smg_rsa.h"
#include "mpi.h"

#include <stdlib.h>
#include <unistd.h>
#include <time.h>
#include <stdio.h>

void err(char *msg)
{
  fprintf(stderr, "%s\n", msg);
  exit(1);
}

void time_entropy_source(int nruns, int noctets) {
	unsigned char buffer[noctets];
	int read, i;
	struct timespec tstart, tend;
	long int diff;

	clock_gettime(CLOCK_MONOTONIC, &tstart);
	for (i=0; i<nruns; i++) {
		read = get_random_octets(noctets,buffer);
		if (read != noctets)
			err("Failed reading from entropy source!");
	}
	clock_gettime(CLOCK_MONOTONIC, &tend);

	diff = tend.tv_sec-tstart.tv_sec;
	double kbps = (nruns*noctets) / (diff*1000.0);
	printf("ENTROPY source timing: %d kB in %ld seconds, at an average speed of %f kB/s over %d runs of %d octets each\n", nruns*noctets, diff, kbps, nruns, noctets);
}

void test_entropy_output(unsigned int noctets, char * filename) {
	FILE * out;
	int source;
	unsigned int nread, total_read, to_read;
	const int buffer_length = 1000;
	unsigned char buffer[buffer_length];

	source = open_entropy_source(ENTROPY_SOURCE);
	if (source <= 0)
		err("unable to access entropy source!");

	out = fopen(filename, "wb");
	if ( !out )
		err("unable to open output file for test_entropy_output!");

	printf("TEST_ENTROPY_SOURCE: reading %u octets from %s ", noctets, ENTROPY_SOURCE);
	total_read = 0;
	while (total_read < noctets) {
		to_read = noctets - total_read;
		if (to_read > buffer_length)
			to_read = buffer_length;

		nread = get_random_octets_from(to_read, buffer, source);
		if (nread > 0) {
			total_read = total_read + nread;
			fwrite(buffer, 1, nread, out);
			fflush(out);
			printf(".");
			fflush(stdout);
		}
	}
	printf("done.\n");

	fclose(out);
	close(source);
}

void test_is_composite(int nruns, char *hex_number, int expected) {
	int i;
	int output;
	int count_ok = 0;
	int source = open_entropy_source(ENTROPY_SOURCE);
	MPI p = mpi_alloc(0);

	mpi_fromstr(p, hex_number);
	printf("TEST is_composite on MPI(hex) ");
	mpi_print(stdout, p, 1);
	for (i=0; i < nruns; i++) {
		printf(".");
		fflush(stdout);
		output = is_composite(p, M_R_ITERATIONS, source);
		if (output == expected)
			count_ok = count_ok + 1;
	}
	printf("done, with %d out of %d correct runs for expected=%d: %s\n", count_ok, nruns, expected, count_ok==nruns? "PASS":"FAIL");
	mpi_free(p);
	close(source);
}

void time_mr(int nruns) {
	struct timespec tstart, tend;
	long int diff;
	int i;
	MPI prime;
	unsigned int noctets = KEY_LENGTH_OCTETS / 2;
	unsigned int nlimbs = mpi_nlimb_hint_from_nbytes(noctets);

	int entropy_source = open_entropy_source(ENTROPY_SOURCE);
	if (entropy_source <= 0)
		err("can't open entropy source!");

	/* first generate a prime of half key length, to make sure M-R will run max number of iterations */
	printf("Generating a prime number of %d octets length for M-R timing test\n", noctets);
	prime = mpi_alloc(nlimbs);
	gen_random_prime(noctets, prime);

	printf("Running timing test for Miller-Rabin with %d repetitions and %d witnesses on prime number ", nruns, M_R_ITERATIONS);
	mpi_print(stdout, prime, 1);
	printf("\n");
	/* now do the actual runs and time it all */
	clock_gettime(CLOCK_MONOTONIC, &tstart);
	for (i=0; i<nruns; i++) {
		if (is_composite(prime, M_R_ITERATIONS, entropy_source))
			printf("FAIL");
		else printf(".");
		fflush(stdout);
	}
	clock_gettime(CLOCK_MONOTONIC, &tend);

	diff = tend.tv_sec-tstart.tv_sec;
	printf("\nTimings on prime number %d octets long, %d runs of MR with %d iterations (witnesses checked) each\n", \
		noctets, nruns, M_R_ITERATIONS);
	printf("Total time: %ld seconds\nTime per MR run: %f seconds\nTime per MR iteration: %f seconds\n",\
			 diff, diff / (1.0*nruns), diff / (1.0*nruns * M_R_ITERATIONS));

	mpi_free(prime);
	close(entropy_source);
}

void test_rpng(int nruns) {
	unsigned int noctets = KEY_LENGTH_OCTETS / 2;
	unsigned int nlimbs = mpi_nlimb_hint_from_nbytes(noctets);
	int entropy_source = open_entropy_source(ENTROPY_SOURCE);
	if (entropy_source <= 0)
		err("can't open entropy source!");

	MPI prime = mpi_alloc(nlimbs);
	int i;

	printf("TEST: random prime number generator with %d runs\n", nruns);
	for (i = 0;i < nruns; i++) {
		gen_random_prime(noctets, prime);
		printf("Run %d: ", i+1);
		mpi_print(stdout, prime, 1);
		if (is_composite(prime, M_R_ITERATIONS, entropy_source))
			printf("  **FAIL**\n");
		else
			printf("  **PASS**\n");
	}

	mpi_free(prime);
	close(entropy_source);
}

void time_rpng(int nruns) {
	struct timespec tstart, tend;
	long int diff;

	unsigned int noctets = KEY_LENGTH_OCTETS / 2;
	unsigned int nlimbs = mpi_nlimb_hint_from_nbytes(noctets);

	int entropy_source = open_entropy_source(ENTROPY_SOURCE);
	if (entropy_source <= 0)
		err("can't open entropy source!");

	MPI prime = mpi_alloc(nlimbs);
	int i;

	printf("TIMING: random prime number generator with %d runs\n", nruns);
	clock_gettime(CLOCK_MONOTONIC, &tstart);
	for (i = 0;i < nruns; i++) {
		gen_random_prime(noctets, prime);
	}
	clock_gettime(CLOCK_MONOTONIC, &tend);

	diff = tend.tv_sec-tstart.tv_sec;

	printf("TOTAL: %ld seconds\n", diff);
	printf("Average: %f seconds to generate one random prime of %d octets length\n", diff / (1.0*nruns), noctets);
	mpi_free(prime);
	close(entropy_source);
}

/* Test encryption+decryption on noctets of random data, using sk
 * Output is written to file.
 */
void test_rsa_keys( RSA_secret_key *sk, unsigned int noctets, FILE *file ) {
	RSA_public_key pk;
	MPI test = mpi_alloc ( mpi_nlimb_hint_from_nbytes (noctets) );
	MPI out1 = mpi_alloc ( mpi_nlimb_hint_from_nbytes (noctets) );
	MPI out2 = mpi_alloc ( mpi_nlimb_hint_from_nbytes (noctets) );

	pk.n = mpi_copy(sk->n);
	pk.e = mpi_copy(sk->e);
	unsigned char *p;
	p = xmalloc(noctets);

	fprintf(file, "TEST encrypt/decrypt on %d octets of random data\n", noctets);
	fflush(file);
	if (get_random_octets( noctets, p) == noctets) {
		mpi_set_buffer( test, p, noctets, 0 );

		fprintf(file, "TEST data:\n");
		mpi_print(file, test, 1);
		fprintf(file, "\n");
		fflush(file);

		public_rsa( out1, test, &pk );
		secret_rsa( out2, out1, sk );

		fprintf(file, "ENCRYPTED with PUBLIC key data:\n");
		mpi_print(file, out1, 1);
		fprintf(file, "\n");
		fflush(file);

		fprintf(file, "DECRYPTED with SECRET key:\n");
		mpi_print(file, out2, 1);
		fprintf(file, "\n");
		fflush(file);

		if( mpi_cmp( test, out2 ) )
			fprintf(file, "FAILED: RSA operation: public(secret) failed\n");
		else
			fprintf(file, "PASSED: RSA operation: public(secret) passed\n");
		fflush(file);

		secret_rsa( out1, test, sk );
		public_rsa( out2, out1, &pk );
		if( mpi_cmp( test, out2 ) )
			fprintf(file, "FAILED: RSA operation: secret(public) failed\n");
		else
			fprintf(file, "PASSED: RSA operation: secret(public) passed\n");
	}
	else
		fprintf(file, "FAILED: not enough bits returned from entropy source\n");

	fflush(file);
	xfree(p);
	mpi_free( pk.n);
	mpi_free( pk.e);

	mpi_free( test );
	mpi_free( out1 );
	mpi_free( out2 );
}

void test_rsa( int nruns, FILE *fkeys, FILE *fout) {
	RSA_secret_key sk;
	int noctets = KEY_LENGTH_OCTETS;
	int noctets_pq = noctets / 2;
	int nlimbs = mpi_nlimb_hint_from_nbytes(noctets);
	int nlimbs_pq = mpi_nlimb_hint_from_nbytes(noctets_pq);
	int i;

	sk.n = mpi_alloc(nlimbs);
	sk.e = mpi_alloc(nlimbs);
	sk.d = mpi_alloc(nlimbs);
	sk.p = mpi_alloc(nlimbs_pq);
	sk.q = mpi_alloc(nlimbs_pq);
	sk.u = mpi_alloc(nlimbs_pq);

	printf("TEST RSA key generation and use with %d runs\n", nruns);
	fflush(stdout);

	for (i = 0;i < nruns; i++) {
		gen_keypair(&sk);
		printf(".");
		fflush(stdout);

		mpi_print(fkeys, sk.n, 1);
		fwrite("\n", sizeof(char), 1, fkeys);

		mpi_print(fkeys, sk.e, 1);
		fwrite("\n", sizeof(char), 1, fkeys);

		mpi_print(fkeys, sk.d, 1);
		fwrite("\n", sizeof(char), 1, fkeys);

		mpi_print(fkeys, sk.p, 1);
		fwrite("\n", sizeof(char), 1, fkeys);

		mpi_print(fkeys, sk.q, 1);
		fwrite("\n", sizeof(char), 1, fkeys);

		mpi_print(fkeys, sk.u, 1);
		fwrite("\n", sizeof(char), 1, fkeys);

		test_rsa_keys(&sk, noctets_pq, fout);
		printf("*");
		fflush(stdout);
	}

	mpi_free(sk.n);
	mpi_free(sk.e);
	mpi_free(sk.d);
	mpi_free(sk.p);
	mpi_free(sk.q);
	mpi_free(sk.u);

}

void test_rsa_exp() {
	MPI msg = mpi_alloc(0);
	MPI expected = mpi_alloc(0);
	MPI result;

	RSA_public_key pk;
	pk.n = mpi_alloc(0);
	pk.e = mpi_alloc(0);

	printf("TEST verify of rsa exponentiation on input data: \n");

	mpi_fromstr(msg, "0x\
5B6A8A0ACF4F4DB3F82EAC2D20255E4DF3E4B7C799603210766F26EF87C8980E737579\
EC08E6505A51D19654C26D806BAF1B62F9C032E0B13D02AF99F7313BFCFD68DA46836E\
CA529D7360948550F982C6476C054A97FD01635AB44BFBDBE2A90BE06F7984AC8534C3\
8613747F340C18176E6D5F0C10246A2FCE3A668EACB6165C2052497CA2EE483F4FD8D0\
6A9911BD97E9B6720521D872BD08FF8DA11A1B8DB147F252E4E69AE6201D3B374B171D\
F445EF2BF509D468FD57CEB5840349B14C6E2AAA194D9531D238B85B8F0DD352D1E596\
71539B429849E5D965E438BF9EFFC338DF9AADF304C4130D5A05E006ED855F37A06242\
28097EF92F6E78CAE0CB97");

	mpi_fromstr(expected, "0x\
1FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF\
FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF\
FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF\
FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF\
FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF003051300\
D0609608648016503040203050004406255509399A3AF322C486C770C5F7F6E05E18FC\
3E2219A03CA56C7501426A597187468B2F71B4A198C807171B73D0E7DBC3EEF6EA6AFF\
693DE58E18FF84395BE");
	result = mpi_alloc( mpi_get_nlimbs(expected) );

	mpi_fromstr(pk.n, "0x\
CDD49A674BAF76D3B73E25BC6DF66EF3ABEDDCA461D3CCB6416793E3437C7806562694\
73C2212D5FD5EED17AA067FEC001D8E76EC901EDEDF960304F891BD3CAD7F9A335D1A2\
EC37EABEFF3FBE6D3C726DC68E599EBFE5456EF19813398CD7D548D746A30AA47D4293\
968BFBAFCBF65A90DFFC87816FEE2A01E1DC699F4DDABB84965514C0D909D54FDA7062\
A2037B50B771C153D5429BA4BA335EAB840F9551E9CD9DF8BB4A6DC3ED1318FF3969F7\
B99D9FB90CAB968813F8AD4F9A069C9639A74D70A659C69C29692567CE863B88E191CC\
9535B91B417D0AF14BE09C78B53AF9C5F494BCF2C60349FFA93C81E817AC682F0055A6\
07BB56D6A281C1A04CEFE1");

	mpi_fromstr( pk.e, "0x10001");

	mpi_print( stdout, msg, 1);
	printf("\n");

	public_rsa( result, msg, &pk);
	if ( mpi_cmp( result, expected) != 0 )
		printf( "FAILED\n");
	else
		printf( "PASSED\n");

	printf("Expected:\n");
	mpi_print( stdout, expected, 1);
	printf("\n");

	printf("Obtained:\n");
	mpi_print( stdout, result, 1);
	printf("\n");

	mpi_free( pk.n );
	mpi_free( pk.e );
	mpi_free( msg );
	mpi_free( expected );
	mpi_free( result );
}

void time_rsa_gen( int nruns ) {
	struct timespec tstart, tend;
	long int diff;
	int i;

	RSA_secret_key sk;
	int noctets = KEY_LENGTH_OCTETS;
	int noctets_pq = noctets / 2;
	int nlimbs = mpi_nlimb_hint_from_nbytes(noctets);
	int nlimbs_pq = mpi_nlimb_hint_from_nbytes(noctets_pq);
	sk.n = mpi_alloc(nlimbs);
	sk.e = mpi_alloc(nlimbs);
	sk.d = mpi_alloc(nlimbs);
	sk.p = mpi_alloc(nlimbs_pq);
	sk.q = mpi_alloc(nlimbs_pq);
	sk.u = mpi_alloc(nlimbs_pq);

	clock_gettime(CLOCK_MONOTONIC, &tstart);
	for (i = 0;i < nruns; i++) {
		gen_keypair(&sk);
	}
	clock_gettime(CLOCK_MONOTONIC, &tend);

	diff = tend.tv_sec-tstart.tv_sec;

	printf("TOTAL: %ld seconds for generating %d key pairs\n", diff, nruns);
	printf("Average (%d runs): %f seconds per TMSR RSA key pair.\n",
				nruns, diff / (1.0*nruns));
	mpi_free(sk.n);
	mpi_free(sk.e);
	mpi_free(sk.d);
	mpi_free(sk.p);
	mpi_free(sk.q);
	mpi_free(sk.u);
}

int main(int ac, char **av)
{
	int nruns;
	int id;
	FILE *fk;
	FILE *fout;

	if (ac<2) {
		printf("Usage: %s number_of_runs/octets [testID]\n", av[0]);
		return -1;
	}
	nruns = atoi(av[1]);

	if (ac < 3)
		id = -1;
	else
		id = atoi(av[2]);

	switch ( id ) {
		case 0:
			printf("Timing entropy source...\n");
			time_entropy_source(nruns, 4096);
			break;
		case 1:
			test_entropy_output(nruns, "entropy_source_output.txt");
			break;
		case 2:
			/* tests on miller-rabin */
			/* a few primes (decimal): 65537, 116447, 411949103, 20943302231 */
			test_is_composite(nruns, "0x10001", 0);
			test_is_composite(nruns, "0x1C6DF", 0);
			test_is_composite(nruns, "0x188DD82F", 0);
			test_is_composite(nruns, "0x4E0516E57", 0);
			/* a few mersenne primes (decimal): 2^13 - 1 = 8191, 2^17 - 1 = 131071, 2^31 - 1 = 2147483647 */
			test_is_composite(nruns, "0x1FFF", 0);
			test_is_composite(nruns, "0x1FFFF", 0);
			test_is_composite(nruns, "0x7FFFFFFF", 0);
			/* a few carmichael numbers, in decimal: 561, 60977817398996785 */
			test_is_composite(nruns, "0x231", 1);
			test_is_composite(nruns, "0xD8A300793EEF31", 1);
			/* an even number */
			test_is_composite(nruns, "0x15A9E672864B1E", 1);
			/* a phuctor-found non-prime public exponent: 170141183460469231731687303715884105731 */
			test_is_composite(nruns, "0x80000000000000000000000000000003", 1);
			break;
		case 3:
			time_mr(nruns);
			break;
		case 4:
			test_rpng(nruns);
			break;
		case 5:
			time_rpng(nruns);
			break;
		case 6:
			fk = fopen("keys.asc", "a");
			if ( fk == NULL )
				err("Failed to open file keys.asc!");
			fout = fopen("check_keys.asc", "a");
			if ( fout == NULL ) {
				fclose(fk);
				err("Failed to open file keys_check.asc!");
			}
			test_rsa(nruns, fk, fout);
			fclose(fk);
			fclose(fout);
			break;
		case 7:
			test_rsa_exp();
			break;
		case 8:
			time_rsa_gen(nruns);
			break;
		default:
			printf("Current test ids:\n");
			printf("0 for timing entropy source\n");
			printf("1 for entropy output test\n");
			printf("2 for is_composite (Miller-Rabin) test\n");
			printf("3 for timing Miller-Rabin\n");
			printf("4 for random prime number generator test\n");
			printf("5 for timing random prime number generator\n");
			printf("6 for testing rsa key pair generation and use; \
writes to keys.asc and check_keys.asc\n");
			printf("7 for testing rsa exponentiation (fixed data)\n");
			printf("8 for timing rsa key pair generator\n");
	}

  return 0;
}