#include "lib3.h"
#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;
}

int crack_aes_ecb(char **plaintext, int blocksize, int offset)
{
	int prepend_data = offset ? 1 : 0;
	int fill_in = offset % blocksize;
	if(fill_in) {
		fill_in = blocksize-fill_in;
		offset += fill_in;
	}
	printf("offset must now be multiple of 16: %i\n", offset % blocksize);
	char *attacker_data = malloc(17+fill_in);
	int length_plaintext;
	char **ciphertexts= malloc(sizeof(char*)*blocksize);
	int *ciphertext_length = malloc(sizeof(char*)*blocksize);
	char *cipher_block = malloc(blocksize+1);
	char *prefix = malloc(blocksize+1);
	int i,j;
	// encrypt text prepended with A, AA, AAA, ...
	memset(attacker_data, 0, 17+fill_in);
	for(i=0;i<blocksize;i++) {
		// fill it with our controlled input
		memset(attacker_data, 0x54, i+fill_in);
		// encrypt the text 
		ciphertext_length[i] = challenge12_and_14_oracle(attacker_data, i+fill_in,&ciphertexts[i], prepend_data);
		//aes_ecb(to_encrypt_block, (length_plaintext+17), ciphertexts[i], key, 16, 1);
	}
	length_plaintext = ciphertext_length[0];
	// we have all ciphertexts now to crack the whole plaintext
	// let's go! 
	*plaintext = malloc(length_plaintext);

	memset(*plaintext, 0, length_plaintext);
	// for every block 
	for(i=0;i<=(length_plaintext-offset)/blocksize;i++) {
		for(j=blocksize-1;j>=0;j--) {
			if(i==0)
				memset(prefix, 0x54, blocksize);
			else
				memcpy(prefix, &((*plaintext)[((i-1)*blocksize)+blocksize-j]), blocksize);

			memcpy(&prefix[j], &((*plaintext)[i*blocksize]), blocksize-j);
			memcpy(cipher_block, &(ciphertexts[j][i*blocksize+offset]), blocksize);
			(*plaintext)[i*blocksize+blocksize-1-j] = create_dictionary_and_match(prefix, cipher_block, key, blocksize);
		}
	}
}

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;
		}
	}
}


int challenge12_and_14_oracle(char *attacker_data, int attacker_data_length, char **ciphertext, int prepend_data)
{
	// target data
	char *base64_task_string = "Um9sbGluJyBpbiBteSA1LjAKV2l0aCBteSByYWctdG9wIGRvd24gc28gbXkgaGFpciBjYW4gYmxvdwpUaGUgZ2lybGllcyBvbiBzdGFuZGJ5IHdhdmluZyBqdXN0IHRvIHNheSBoaQpEaWQgeW91IHN0b3A/IE5vLCBJIGp1c3QgZHJvdmUgYnkK";

	char *target_data= malloc(strlen(base64_task_string));
	// unbases it
	int length_cleartext = decode_base64(base64_task_string, target_data);
	
	static char *prepend = "";
	static int prepend_length = 0;
	
	// generate random prefix only the first time
	if(prepend_length == 0 && prepend_data) {
		prepend_length = random_number_between(17,200);
		printf("prepend %i bytes (%i full blocks)\n", prepend_length, prepend_length/16);
		prepend = malloc(prepend_length);
		generate_random_bytes(prepend, prepend_length);
	}

	//target_data = "";

	// mix to one big string and encrypt
	int plaintext_length = prepend_length+attacker_data_length+strlen(target_data);
	char *plaintext = malloc(plaintext_length);
	*ciphertext = malloc(plaintext_length);
	memcpy(plaintext, prepend, prepend_length);
	memcpy(&plaintext[prepend_length], attacker_data, attacker_data_length);
	memcpy(&plaintext[prepend_length+attacker_data_length], target_data, strlen(target_data));

	aes_ecb(plaintext, plaintext_length, *ciphertext, key, 16, 1);
	return plaintext_length;
}

#define BLOCKSIZE 16
int aes_ecb_detect_prepended_data()
{
	int i;
	int prepend_bytes;
	char *ciphertext_round;
	char *ciphertext_without;
	char a_block[2*BLOCKSIZE+1];
	char a_block_enc[BLOCKSIZE+1];

	memset(a_block, 'A', 2*BLOCKSIZE);
	a_block[2*BLOCKSIZE+1] = '\0';

	int number_equal_blocks = 0;
	int ciphertext_without_len = challenge12_and_14_oracle(a_block,0, &ciphertext_without, 1);
	int ciphertext_round_len = challenge12_and_14_oracle(a_block, 2*BLOCKSIZE, &ciphertext_round, 1);

	// count equal blocks
	for(i=0;i<(ciphertext_round_len/BLOCKSIZE);i++) {
		if(memcmp(&ciphertext_without[i*BLOCKSIZE], &ciphertext_round[i*BLOCKSIZE], BLOCKSIZE))
			break;
	}
	prepend_bytes = i*BLOCKSIZE;
	aes_ecb(a_block, BLOCKSIZE, a_block_enc, key, BLOCKSIZE, 1);	
	for(i=BLOCKSIZE*2;i>0;i--) {
		a_block[i] = '\0';
		challenge12_and_14_oracle(a_block, i, &ciphertext_round, 1);	
		if(memcmp(a_block_enc, &ciphertext_round[prepend_bytes+16], BLOCKSIZE))
			break;
	}

	prepend_bytes += (2*BLOCKSIZE-i-1);
	return prepend_bytes;
}

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, int cbc_mode)
{
	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
	if(cbc_mode) {
		unencrypted = pkcs7_padding(res, strlen(res), 16);
		*encrypted = malloc(strlen(unencrypted));
		aes_cbc(unencrypted, strlen(unencrypted), *encrypted, key, iv , 1);
		return strlen(unencrypted);
	}
	// otherwise is CTR mode
	else {
		*encrypted = malloc(strlen(res));
		aes_ctr(res, strlen(res), *encrypted, key, iv);
		return strlen(res);
	}
}

void challenge16_decrypt(char *encrypted, int length, int cbc_mode)
{
	char *unencrypted = malloc(length);
	char *unpadd= malloc(length);

	if(cbc_mode) {
		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);
	}
	// we are in ctr mode
	else {
		aes_ctr(encrypted, length, unencrypted, key, iv);
		printf("unencrpted string: %s\n", unencrypted);
	}
}