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


void generate_random_printable_bytes(char *s, unsigned int length)
{
	int i;

	for(i=0;i<length;i++)
		s[i] = random_number_between(33,127);
}

void format_request(char **res, char *content, char *secret)
{
	char *request_template =
		"POST / HTTP/1.1\n"
		"Host: hapless.com\n"
		"Cookie: sessionid=%s\n"
		"Content-Length: %i\n\n\n"
		"%s";

	// 10 for the content length as max
	*res = malloc(strlen(request_template)+10+strlen(content)+strlen(secret));
	sprintf(*res, request_template, secret, strlen(content), content);
}

void my_compress(char *toCompress, char **compressed, int *compressed_length)
{
	// 20 bytes for metadata and bla but normally it should be smaller
	// than the uncompressed one
	*compressed = malloc(strlen(toCompress)+20);
	// compress
	z_stream defstream;
	defstream.zalloc = Z_NULL;
	defstream.zfree = Z_NULL;
	defstream.opaque = Z_NULL;

	defstream.avail_in = strlen(toCompress);
	defstream.next_in = toCompress;
	defstream.avail_out = strlen(toCompress)+20;
	defstream.next_out = *compressed;

	deflateInit(&defstream, Z_BEST_COMPRESSION);
	deflate(&defstream, Z_FINISH);
	deflateEnd(&defstream);

	*compressed_length = defstream.total_out;
	//printf("uncompressed/compressed: %i/%i\n", strlen(toCompress), defstream.total_out);

}
int oracle_count = 0;

int compression_oracle(char *content, char* secret, int stream_cipher)
{
	oracle_count++;
	char *res;
	char *compressed;
	int compressed_length;
	format_request(&res, content, secret);
	//printf("%s\n", res);
	my_compress(res, &compressed, &compressed_length);
	// encrypt
	// using AES_CTR, plaintext and ciphertext length are the same
	// do it anyway
	if(stream_cipher == 1) {
		char *encrypted = malloc(compressed_length);
		generate_random_bytes(key, 16);
		generate_random_bytes(nonce, 16);
		aes_ctr(compressed, compressed_length, encrypted, key, nonce);
		return compressed_length;
	}
	else {
		// do CBC encryption aka add padding
		char *encrypted;
		int pad_len = aes_cbc_padded(compressed, compressed_length, &encrypted,
				key, iv, 1);
		return pad_len;
	}
}

int max_int_index(int *numbers, int length)
{
	int i;
	int max = 0;
	int index = 0;

	for(i=0;i<length;i++) {
		if(numbers[i] > max) {
			max = numbers[i];
			index = i;
			}
	}

	return index;
}

int main()
{
	srand(time(NULL));
	char *known_text = "Cookie: sessionid=";
	char *secret = "TmV2ZXIgcmV2ZWFsIHRoZSBXdS1UYW5nIFNlY3JldCE=";
	// generate password between 33, 127 
	// assume here that the session id is printable (HTTP is text protocol)
	char *t = malloc(strlen(known_text)+50);
	memset(t, 0, strlen(known_text)+50);
	memcpy(t, known_text, strlen(known_text));
	int max, res, i, k;
	
	// assume that the string we wanne find is less than 100 bytes
	for(k=0;k<100;k++) {
		t[strlen(known_text)+k] = 0x09;
		max = compression_oracle(t, secret, 1);
		for(i=33;i<128;i++) {
			t[strlen(known_text)+k] = (char) i;			
			res = compression_oracle(t, secret, 1);
			if(res < max)
				break;
		}
		if(max == res)
			break;
	}

	printf("found string : %s\n", t);
	printf("should string: Cookie: sessionid=%s\n", secret);
	// doing binary search above would reduce the calls to the orcale
	// significant
	printf("calls to the oracle: %i\n", oracle_count);
	
	
	// with cbc
	// length will always be a multiple of the block size
	oracle_count = 0;
	printf("with CBC mode\n");
	// because we only get multiple of 16 as result we have to take
	// care that a complete block gets compressed
	// we than need control over the size of the request so that
	// is is a multiple of 16 BEFORE padding
	// otherwise we compress 15 bytes and use one block less
	char *kt = malloc(strlen(known_text)+strlen(secret));
	memcpy(kt, known_text, strlen(known_text));
	memcpy(&kt[strlen(known_text)], secret, strlen(secret));
	// one block which should be compressed away
	// 6 is the alignment to get a full block
	char rand[17+6];
	char ttt[18+8];
	int j,l;
	generate_random_printable_bytes(rand, 16+8);
	rand[17] = '\0';
	char *raw_request;
	char *compressed;
	int compressed_length;
	int pad;
	int counts[127-33];
	int times = 8;
	char found_sec_cbc[100];

	format_request(&raw_request, "AAAAAAAAAAAAAAAA", secret);
	my_compress(raw_request, &compressed, &compressed_length);
	__pkcs7_padding(compressed, compressed_length, 16, &pad);
	int max_ = compression_oracle(rand, secret, 0);

	memcpy(&ttt[pad], &known_text[strlen(known_text)-16], 16);
	for(l=0;l<strlen(secret);l++) {
		memset(&counts, 0, sizeof(int)*(127-33));
		generate_random_printable_bytes(ttt, pad);
		for(i=0;i<15;i++)
			ttt[pad+i] = ttt[pad+i+1];
		for(k=0;k<10;k++) {
			for(i=33;i<127;i++) {
				ttt[pad+15] = i;
				ttt[pad+16] = '\0';
				int len = compression_oracle(ttt, secret, 0);
				if( len < max_)
					break;
			}
			counts[i-33]++;
		}
		i = max_int_index(counts, 127-33);
		found_sec_cbc[l] = (char) (i + 33);
	}
	found_sec_cbc[l] = '\0';
	printf("should: %s\n", secret);
	printf("is    : %s\n", found_sec_cbc);
	printf("calls to the oracle: %i\n", oracle_count);
}