#include "../lib/lib.h"
#include "../lib/lib2.h"
#include "../lib/lib3.h"
#include "../lib/lib4.h"
#include "../lib/lib5.h"
#include "../lib/lib7.h"

/**
  * MD-Hash implementation with AES-128. Outputs only 2^16 bits of the hash
  **/ 
void  md_hash_aes_128(char *input, int input_len, char *state, char *hash, int bytes, int print)
{
	int blocksize = 16;
	int i;
	int pad_len;
	char h[blocksize];
	char key[blocksize];
	char h_hex[32];
	char *input_padded = __pkcs7_padding(input, input_len, 16, &pad_len);

	// set the intial state
	memset(h, 0, blocksize);
	memcpy(h, state, bytes);
	if(print == 1)
		printf("input len is: %i, in blocks: %i\n", input_len, input_len/blocksize);

	for(i=0;i<input_len/blocksize;i++) {
		// only use the first bytes of the hash as state
		// zero the others out
		memset(key, 0, 16);
		memcpy(key, h, bytes);
		if(print == 1) {
			hex_binary_to_string(h, h_hex, 16);
			printf("%i: key is: 5s\n", i, h_hex);
			memset(key, 0, 16);
			memcpy(key, h, bytes);
		}
		aes_ecb(&input[i*blocksize], 16, h, key, blocksize, 1);
	}
	
	// output only bytes bytes
	memcpy(hash, h, bytes);
	free(input_padded);
}


void cascaded_md_hash_functions(char *input, int input_len, char *hash)
{
	// F: cheap one: 2 bytes
	// G: expensive: 16 bytes
	char state[16];
	memset(state, 0, 16);
	md_hash_aes_128(input, input_len, state, hash, 2, 0);
	md_hash_aes_128(input, input_len, state, &hash[2], 4, 0);
}

void add_one(unsigned char *s, int s_len)
{
	if(s[s_len-1] == 255) {
		s[s_len-1] = 0x00;
		add_one(s, s_len-1);
	}
	else {
		s[s_len-1] += 1;
	}
}

int char_array_is_max(unsigned char *s, int s_len)
{
	int i;

	for(i=0;i<s_len;i++)
		if(s[i] != 255)
			return 0;

	return 1;
}

int _2_collision_finding_machine(char *state, char *in1, char *in2, int bytes)
{
	// find a collsion for the given state
	char hash1[16];
	char hash2[16];
	char op_state[16];

	memset(in1, 0, 16);
	memset(in2, 0, 16);

	do {
		if(char_array_is_max(in2, 16) == 1) {
			add_one(in1, 16);
			memset(in2, 0, 16);
		}
		add_one(in2, 16);
		memcpy(op_state, state, 16);
		md_hash_aes_128(in1, 16, state, hash1, bytes, 0);
		memcpy(op_state, state, 16);
		md_hash_aes_128(in2, 16, state, hash2, bytes, 0);
	} while(memcmp(hash1, hash2, bytes) != 0 && (memcmp(in1, in2, 16) != 0));
	memcpy(state, hash1, bytes);
}


char **collisions = NULL;
unsigned int coll_numbers = 0;

void generate_collsions(char *prefix, int pref_len, char *coll1, char *coll2, int length)
{
	if(length == 0) {
		collisions[coll_numbers++] = prefix;
		return;
	}
		
	char *print = malloc(pref_len+16);
	memcpy(print, prefix, pref_len);
	memcpy(&print[pref_len], coll1, 16);
	generate_collsions(print, pref_len+16, &coll1[16], &coll2[16], length-1);
	char *print2 = malloc(pref_len+16);
	memcpy(print2, prefix, pref_len);
	memcpy(&print2[pref_len], coll2, 16);
	generate_collsions(print2, pref_len+16, &coll1[16], &coll2[16], length-1);
	//free(print);
}

/**
  * This functions computes 2^n collision for a MD-Hash function given
  * an 2 collision finding machine
  **/
int f(int n, int bytes)
{
	int i;
	char state[16];
	char *block1 = malloc(n*16);
	char *block2 = malloc(n*16);
	char *hex_block = malloc(n*32);
	memset(state, 0, 16);

	printf("find collision blocks..\n");
	for(i=0;i<n;i++) {
		_2_collision_finding_machine(state, &block1[i*16], &block2[i*16], bytes);
	}
	printf("generating collsions out of the blocks\n");
	generate_collsions("", 0, block1, block2, n);
	free(block1);
	free(block2);
}

void print_collisions(int coll_len, int bytes)
{
	int i;
        char *coll_hex = malloc(coll_len*2);
        char hex_hash[32];
        char hash[16];
        char state[16];

	for(i=0;i<coll_numbers;i++) { 
		printf("processing collision at %p\n", collisions[i]);
		memset(state,0,16);
		md_hash_aes_128(collisions[i], coll_len, state, hash, bytes, 0);
		hex_binary_to_string(collisions[i], coll_hex, coll_len);
		hex_binary_to_string(hash, hex_hash, bytes);
		printf("data: %s:\t hash: %s\n", coll_hex, hex_hash);
	}
//	free(coll_hex);
}

/**
  * find a collsion under both of the cascaded hash function
  *
  **/
void f_cascaded()
{
	// generate 2^b/2 collsions with the cheap one
	int i,j;
	char state[16];
	int found_collisions_expensive = 0;
	int bytes_f = 2;
	int bytes_g = 3;
	int __collisions = 1;
	int coll_controller = 12;
	for(i=0;i<coll_controller;i++)
		__collisions *= 2;

	collisions = malloc(sizeof(char *)*__collisions);
	printf("generate %i collisions\n", __collisions);
	f(coll_controller, bytes_f);
	char **hashes = malloc(sizeof(char *)*__collisions);
	// check if one also collide in the not so cheap one
	printf("check for collisions with the expensive hash function...\n");
	for(i=0;i<coll_numbers;i++) {
		hashes[i] = malloc(bytes_g);
		memset(state, 0, 16);
		md_hash_aes_128(collisions[i], coll_controller*16, state, hashes[i], bytes_g, 0);

		for(j=0;j<i;j++) {
			if(i != j && memcmp(hashes[j], hashes[i], bytes_g) == 0) {
				char hex_hash[32];
				char *coll_hex = malloc(coll_controller*32);
				printf("collision in expensive hash function found: %i, %i ", i, j);
				hex_binary_to_string(hashes[i], hex_hash, bytes_g);
				printf("with hash: %s\n", hex_hash);

				hex_binary_to_string(collisions[i], coll_hex, coll_controller*16);
				printf("data %i: %s\n", i, coll_hex);
				hex_binary_to_string(collisions[j], coll_hex, coll_controller*16);
				printf("data %i: %s\n", j, coll_hex);
				found_collisions_expensive++;
			}
		}
	}
	printf("found %i collisions in the expensive one out of %i collision candiate\n", found_collisions_expensive, coll_numbers);
}


int part1()
{
	int i;
	int n = 3;
	int  bytes = 2;
	int __coll_numbers = 2;
	for(i=1;i<n;i++)
		__coll_numbers = __coll_numbers * 2;
	printf("2^n is: %i\n", __coll_numbers);
	collisions = malloc(sizeof(char *)*__coll_numbers);

	f(n, bytes);
	printf("found collision: %i\n", coll_numbers);
	print_collisions(n*16, bytes);

}


int part2()
{
	if(collisions != NULL)
		free(collisions);
	coll_numbers = 0;
	f_cascaded();
}

int main()
{
	srand(time(NULL));

	part1();
	printf("\n\npart2:\n\n");
	part2();
}