#include "lib2.h"
#include "lib.h"

/**
 * appends PKCS#7 padding to string. devide string in blocks of size blocksize_bytes
 * and append to last block so that it is also of blocksize length
 */

char *__pkcs7_padding(char *string, int length_string, int blocksize, int *padding)
{
	char *result = NULL;
	int i;
	int value = blocksize - (length_string % blocksize);
	result = malloc(length_string+value+1);
	memcpy(result, string, length_string);

	for(i=0;i<value;i++) {
		result[length_string+i] = (char) value;
	}
	result[length_string+i] = '\0';
	*padding = value;
	return result;
}

char *pkcs7_padding(char *string, int length_string, int blocksize)
{
	int not_interested;
	return __pkcs7_padding(string, length_string, blocksize, &not_interested);
}
/**
 * unpadd a string
 * @param in string which should be unpadded
 * @param length_in length of paramter length_in
 * @param unpadded place where the unpadded text should be stored
 * @param blocksize size of the block to which should be padded
 */
int valid_pkcs7_padding(const char *in, int length_in, char *unpadded, int blocksize)
{
	int i, padding_length;
	// look how many equal bytes are at the end
	// ignore last byte \0
	for(i=length_in-1;i>0;i--) {
		if(in[i] != in[i-1])
			break;
	}
	
	padding_length = length_in - i;
	if ((length_in % blocksize) != 0)
		return 0;

	if(in[length_in-1] != padding_length)
		return 0;

	memcpy(unpadded, in, (length_in-padding_length));
	unpadded[length_in-padding_length+1] ='\0';
	return 1;
}

/**
 * decrypts content which is encrypted in AES CBC mode
 * @param in input content
 * @param length_in length of parametere in
 * @param out place where the decrypted content should be written to
 * @param string_key key with which the content in in has been decrypted
 * @param iv initalization vector
 */
int aes_cbc(char *in, int length_in, char *out, unsigned char *string_key, char *init_vector, int encrypt)
{
	char iv[16];
	AES_KEY key;
	int number_blocks = length_in / 16;
	int i, j;
	unsigned char ciphertext[128+1];
	unsigned char tmp_after_aes[128+1];
	unsigned char cleartext[128+1];
	// set the key and bits
	if(encrypt)
		AES_set_encrypt_key(string_key, 128, &key);
	else
		AES_set_decrypt_key(string_key, 128, &key);

	memcpy(init_vector, iv, 16);

	// implement cbc mode
	for(i=0;i<number_blocks;i++) {
		if (!encrypt) {
			//do aes decryption
			AES_decrypt(&in[i*16], tmp_after_aes, &key);
			// xor
			xor_string(iv, tmp_after_aes, &out[i*16], 16, 16);
			// this ciphertext block is the next iv
			for(j=0;j<16;j++) {
				iv[j] = in[i*16+j];
			}
		}
		else {
			// first xor
			xor_string(iv, &in[i*16], tmp_after_aes, 16, 16);
			// aes encrypt
			AES_encrypt(tmp_after_aes, &out[i*16], &key);
			// ciphertext is next iv
			for(j=0;j<16;j++) {
				iv[j] = out[i*16+j];
			}
		}
	}
	return 0;

}

int aes_ecb(char *in, int length_in, char *out, unsigned char *string_key,
		int blocksize_bytes, int encrypt)
{
	AES_KEY key;
	int number_blocks = length_in / blocksize_bytes;
	int i;

	if(encrypt)
		AES_set_encrypt_key(string_key, blocksize_bytes*8, &key);
	else
		AES_set_decrypt_key(string_key, blocksize_bytes*8, &key);
	
	for(i=0;i<number_blocks;i++) {
		if(encrypt) {
			AES_encrypt(&in[i*blocksize_bytes], &out[i*blocksize_bytes],&key);
		}
		else {
			AES_decrypt(&in[i*blocksize_bytes], &out[i*blocksize_bytes],&key);
		}
	}
}



int random_number_between(int min, int max)
{
	return (rand() % (max-min) + min);
}


int generate_random_bytes(char *buf, int length_key_bytes)
{
	int i;
	for(i=0;i<length_key_bytes;i++) {
		buf[i] = (char) random_number_between(0,255);
	}

}

char *encrypt_with_random_bytes(char *toencrypt, int length, int ecb)
{
	int toappend_before = random_number_between(5,10);
	int toappend_after= random_number_between(5,10);

	char random_bytes[11];
	char key[17];
	// blocksize is 16 so allocate engouh
	char *result = malloc(length+toappend_before+toappend_after+17);
	char *ciphertext = malloc(length+toappend_before+toappend_after+17);

	generate_random_bytes(random_bytes, toappend_before);
	
	memcpy(result, random_bytes, toappend_before);
	memcpy(&result[toappend_before], toencrypt, length);

	generate_random_bytes(random_bytes, toappend_after);

	memcpy(&result[length+toappend_before+toappend_after], random_bytes, toappend_after);
	
	generate_random_bytes(key, 16);


	if(ecb)
		aes_ecb(result, (length+toappend_before+toappend_after), ciphertext, key, 16, 1);
	else {
		char iv[16];
		memset(iv, 0, 16);
		aes_cbc(result, (length+toappend_before+toappend_after), ciphertext, key, iv, 1);
	}

	free(result);
	return ciphertext;
}


int detect_blocksize_ecb(char *text, int length_text, char *key)
{
	int MAX_BLOCK_SIZE = 50;
	int i,j;
	
	int breaking = 0;
	char *encrypt_string = malloc(MAX_BLOCK_SIZE + length_text+17);
	char *encrypted = malloc(MAX_BLOCK_SIZE + length_text+17);
	char last_first_block[MAX_BLOCK_SIZE];
	memset(last_first_block, 0, MAX_BLOCK_SIZE);

	for(i=1;i<MAX_BLOCK_SIZE;i++) {
		memset(encrypt_string, 'A', i);
		memcpy(&encrypt_string[i+1], text, length_text);
		// enrypt this and break if the first block did not change
		aes_ecb(encrypt_string, length_text, encrypted, key, 16, 1);
		for(j=0;j<i;j++) {
			if (last_first_block[j] != encrypted[j]) {
				break;
			}
		}
		memcpy(last_first_block,encrypted,i);
		if(j+1==i && i!=1)
			break;
	}

	free(encrypt_string);
	free(encrypted);

	return i-1;
}
/**
  * cracks one AES-ECB block
  */
int crack_aes_ecb(char *text, int length_text, char *plaintext, char *key, int blocksize)
{
	char *to_encrypt_block = malloc(length_text + 17);
	char **ciphertexts= malloc(sizeof(char*)*blocksize);
	char *cipher_block = malloc(blocksize+1);
	char *prefix = malloc(blocksize+1);
	int i,j;

	for(i=0;i<blocksize;i++) {
		ciphertexts[i] = malloc(length_text + 17);
		memcpy(&to_encrypt_block[i],text, length_text);
		// fill it with our controlled input
		memset(to_encrypt_block, 0x54, i);
		// encrypt the block
		aes_ecb(to_encrypt_block, blocksize, ciphertexts[i], key, 16, 1);
	}
	// we have all ciphertexts now to crack the whole plaintext
	// let's go!

	memset(plaintext, 0, length_text);
	// for every block 
	for(i=1;i<2;i++) {
		for(j=blocksize-1;j>=0;j--) {
			if(i==0)
				memset(prefix, 0x54, blocksize);
			else
				memcpy(prefix, &text[((i-1)*blocksize)+blocksize-j], blocksize);
			memcpy(&prefix[j], &text[i*blocksize], blocksize-j);
			printf("prefix:%s\nendprefix\n", prefix);
			memcpy(cipher_block, &(ciphertexts[j][i*blocksize]), blocksize);
			plaintext[i*blocksize+blocksize-1-j] = create_dictionary_and_match(prefix, cipher_block, key, blocksize);
		}
	printf("plaintext so far: %s\n", plaintext);
	}

}

char create_dictionary_and_match(char *prefix, char *match, char *key, int blocksize_bytes)
{
	char *dict_string = malloc(blocksize_bytes);
	char *cipher_block = malloc(blocksize_bytes);
	char *hex_tmp= malloc(blocksize_bytes*2);
	int i;
	hex_binary_to_string(match, hex_tmp, blocksize_bytes);
	//printf("%s\n", hex_tmp);
	memcpy(dict_string, prefix, blocksize_bytes);
	for(i=0;i<255;i++) {
		dict_string[blocksize_bytes-1] = (char) i;
		printf("%s\n", dict_string);
		aes_ecb(dict_string, blocksize_bytes, cipher_block, key, blocksize_bytes, 1);
		hex_binary_to_string(cipher_block, hex_tmp, blocksize_bytes);
	//	printf("%s\n", hex_tmp);
		if(memcmp(cipher_block, match, blocksize_bytes) == 0) {
			//printf("found charatcer: %i\n", i);
			return (char) i;
		}
	}
}

struct key_value_pair *parse_key_value(char *string, int length_string)
{
	char *str1, *str2, *tmp, *tmp2;
	struct key_value_pair *pair = malloc(sizeof(struct key_value_pair));
	char *saveptr1;
	char *saveptr2;

	for(str1 = string; ; str1 = NULL) {
		tmp2 = strtok_r(str1, "&",  &saveptr1);
		if (tmp2 == NULL)
			break;

		for(str2 = tmp2; ; str2 = NULL) {
			tmp = strtok_r(str2, "=", &saveptr2);
			if (tmp == NULL)
				break;

			if(str2 == NULL) {
				pair->value = malloc(strlen(tmp));
				strcpy(pair->value, tmp);
			} else {
				pair->key = malloc(strlen(tmp));
				strcpy(pair->key, tmp);
			}
		}
		printf("found pair: %s, %s\n", pair->key, pair->value);
	}
	return pair;
}

char *__profile_for(char *email)
{
	char *ret;
	char *before = "email=";
	char *after = "&uid=10&role=user";
	
	if(strchr(email, '=') || strchr(email, '&'))
		return NULL;

	ret = malloc(strlen(email) + strlen(before) + strlen(after)+1);
	strncpy(ret, before, strlen(before));
	strncat(ret, email, strlen(email));
	strncat(ret, after, strlen(after));

	printf("%s\n", ret);
	return ret;

}

char *profile_for(char *email)
{
	char *unencrpyted_profile_string = __profile_for(email);
	char *ret = malloc(strlen(unencrpyted_profile_string));

	aes_ecb(unencrpyted_profile_string , strlen(unencrpyted_profile_string),
				ret, key, 16, 1);

	return ret;
}

void send_user(char *encrypted_user, int length)
{
	char *unencrypted_user = malloc(length);
	aes_ecb(encrypted_user, length, unencrypted_user , key, 16, 0);
	printf("Got user: %s\n", unencrypted_user);
	parse_key_value(unencrypted_user, strlen(unencrypted_user));
}


int challenge16_encrypt(char *input, char **encrypted)
{
	char *prepend = "comment1=cooking\%20MCs;userdata=";
	char *append = ";comment2=\%20like\%20a\%20pound\%20of\%20bacon";
	char *unencrypted;


	char printable_equal[3];
	char printable_semicolon[3];
	int i, count;
	int quote_char = 0;

	for(i=0;i<strlen(input);i++)
		if(input[i] == ';' || input[i] == '=')
			quote_char++;

	// = and ; get %3d and %3b
	char *res = malloc(strlen(input-quote_char) + quote_char*3 +
			strlen(prepend) + strlen(append));

	hex_binary_to_string("=", printable_equal, 1);
	hex_binary_to_string(";", printable_semicolon, 1);


	memcpy(res, prepend, strlen(prepend));

	for (i=0, count = strlen(prepend);i<strlen(input);i++, count++) {
		if (input[i] == '=') {
			res[count++] = '%';
			res[count++] = printable_equal[0];
			res[count] = printable_equal[1];
		} else if (input[i] == ';') {
			res[count++] = '%';
			res[count++] = printable_semicolon[0];
			res[count] = printable_semicolon[1];
		}
		else
			res[count] = input[i];
	}

	memcpy(&res[strlen(prepend)+strlen(input)+2*quote_char], append, strlen(append));
	res[strlen(prepend)+strlen(input)+strlen(append)+2*quote_char+1] = '\0';	
	// padding
	unencrypted = pkcs7_padding(res, strlen(res), 16);
	*encrypted = malloc(strlen(unencrypted));
	aes_cbc(unencrypted, strlen(unencrypted), *encrypted, key, iv , 1);
	return strlen(unencrypted);
}

void challenge16_decrypt(char *encrypted, int length)
{
	char *unencrypted = malloc(length);
	char *unpadd= malloc(length);
	
	aes_cbc(encrypted, length, unencrypted, key, iv, 0);
	// unpadd
	int ret = valid_pkcs7_padding(unencrypted, length, unpadd, 16);
	if(!ret) {
		printf("no valid padding!\n");
		return;
	}
	// look for string ;admin=true;
	printf("unencrpyted string: %s\n", unpadd);
}