#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include "flow.h"
#include <openssl/evp.h>
#include <openssl/dh.h>
#include <openssl/bn.h>
#include <openssl/err.h>
#include <openssl/rand.h>
#include <openssl/ssl.h>
#include "slitheen.h"

#define PRE_MASTER_LEN 256

static flow_table *table;

/* Initialize the table of tagged flows */
int init_flow_table(void) {

	table = malloc(sizeof(flow_table));
	
	table->table = (flow *) malloc(sizeof(flow)*MAX_FLOWS);
	if( table->table == NULL){
		fprintf(stderr, "malloc failed.\n");
		return 1;
	}
	table->len = 0;
	table->max_len = MAX_FLOWS;

	return 0;
}


/* Add a new flow to the tagged flow table */
flow *add_flow(flow newFlow) {
	flow *ptr;

	if(table->len == table->max_len){
		//grow_table();
		NULL;
	}
	printf("there are %d flows in the table\n", table->len);

	ptr = table->table + table->len;
	newFlow.state = TLS_CLNT_HELLO;
	newFlow.in_encrypted = 0;
	newFlow.out_encrypted = 0;
	newFlow.packet_chain = NULL;
	newFlow.finish_md_ctx = EVP_MD_CTX_create();
	const EVP_MD *md = EVP_sha384();
	EVP_DigestInit_ex(newFlow.finish_md_ctx, md, NULL);

	*ptr = newFlow;

	table->len ++;

	return ptr;
}

/* Updates the flow state */
int update_flow(flow *f) {
	uint8_t *record;
	const struct record_header *record_hdr;
	const struct handshake_header *handshake_hdr;

	uint8_t *p = f->packet_chain->data;

	record_hdr = (struct record_header*) p;
	int record_len;
	int data_len;


	//printf("record version(major): %d.\n", (record_hdr->version&0xFF00)>>8);
	//printf("record version(minor): %d.\n", record_hdr->version&0xFF);
	//printf("record length: %d.\n", RECORD_LEN(record_hdr));

	record_len = RECORD_LEN(record_hdr)+RECORD_HEADER_LEN;
	data_len = f->packet_chain->data_len;
	packet *current = f->packet_chain;
	int incoming = current->incoming;
	record = calloc(1, record_len);
	for(int i=0; (i<data_len) && (i<record_len); i++){
		record[i] = p[i];
	}
	//printf("record len: %d, data_len: %d\n", record_len, data_len);
	while(record_len > data_len) {
		if(current->next == NULL){
			//printf("Don't have enought to reconstruct record\n");
			free(record);
			return 0;
		}
		if(current->next->seq_num != current->seq_num + current->len){
			printf("Missing packet: seq_num= %d, datalen= %d, nextseq= %d\n", current->seq_num, current->len, current->next->seq_num);
			free(record);
			return 0;
		}

		current = current->next;
		p = current->data;
		int i;
		for(i=0; (i<current->data_len) && (i+data_len < record_len); i++){

			record[data_len] = p[i];
		}
		//printf("Filled %d\n", i);
		data_len += current->data_len;
	}

	switch(record_hdr->type){
		case HS:
			p = record;
			p += RECORD_HEADER_LEN;

			//int size_hs = HANDSHAKE_MESSAGE_LEN(handshake_hdr);
			printf("Handshake Message:\n");

			if((incoming && f->in_encrypted) || (!incoming && f->out_encrypted)){
				decrypt_fin(f, p, record_len - RECORD_HEADER_LEN, incoming);
				p += EVP_GCM_TLS_EXPLICIT_IV_LEN;
			}
			handshake_hdr = (struct handshake_header*) p;
			f->state = handshake_hdr->type;

			/* Now see if there's anything extra to do */
			switch(f->state){
			/* Checks to see if this is a possibly tagged hello msg */
				case TLS_CLNT_HELLO: 
					/* Expecting server hello msg */
					printf("Received client hello!\n");
					//update_finish_hash(f, p);
					break;
				case TLS_SERV_HELLO:
					extract_server_random(f, p);
					//update_finish_hash(f, p);
					printf("Received server hello!\n");
					break;
				case TLS_NEW_SESS:
					printf("Received new session ticket!\n");
					break;
				case TLS_CERT:
					//update_finish_hash(f, p);
					printf("Received certificate!\n");
					break;
				case TLS_SRVR_KEYEX:
					//update_finish_hash(f, p);
					printf("Received server key exchange!\n");
					/* Need to extract server params */
					if(extract_parameters(f, p)){
						printf("Error extracting params\n");
					}
					if(compute_master_secret(f)){
						printf("Error computing master secret\n");
					}
					break;
				case TLS_CERT_REQ:
					//update_finish_hash(f, p);
					printf("Received certificate request!\n");
					break;
				case TLS_SRVR_HELLO_DONE:
					//update_finish_hash(f, p);
					printf("Received server hello done!\n");
					break;
				case TLS_CERT_VERIFY:
					//update_finish_hash(f, p);
					printf("Received certificate verify!\n");
					break;
				case TLS_CLNT_KEYEX:
					//update_finish_hash(f, p);
					printf("Received client key exchange!\n");
					break;
				case TLS_FINISHED:
					printf("Received finished message!\n");
					break;
				default:
					printf("Error? %02x\n",p[0]);
					break;
			}
			break;
		case APP:
			printf("Application Data\n");
			break;
		case CCS:
			printf("Change of Cipher Spec\n");
			if(incoming){
				f->in_encrypted = 1;
			} else {
				f->out_encrypted = 1;
			}
			/*Initialize ciphers */
			init_ciphers(f);
			break;
		case ALERT:
			printf("Alert\n");
			break;
		case HB:
			printf("Heartbeat\n");
			break;
		default:
			printf("Error: Not a Record\n");
			//TODO: later figure this out, for now delete
			f->packet_chain = f->packet_chain->next;
			if( f->packet_chain != NULL){
				update_flow(f);
			}
			return 0;
	}

	if(record_len == data_len){
		/* record ended on packet boundary */
		//printf("record consumed packet\n");
		f->packet_chain = current->next;
	} else {
		/* need to update data */
		f->packet_chain = current; //TODO: make current
		current->data = current->data + (current->data_len - (data_len - record_len));
		current->data_len = data_len - record_len;
		update_flow(f);
	}

	free(record);
	return 0;
}

int remove_flow(int index) {
	int i;
	flow *ptr;

	if(index){
		ptr = table->table + index -1;
		for(i=0; i< table->len - index; i++){
			ptr += i;
			*ptr = *(ptr + 1);
		}
		table->len --;
	} else {
		return 1;
	}

	printf("flow removed!\n");
	return 0;
}

int grow_table() {
	return 0;
}

/** Returns the index of a flow in the flow table if
 *  it exists, returns 0 if it is not present.
 */
int check_flow(flow observed){
	/* Loop through flows in table and see if it exists */
	int i;
	flow *candidate = table->table;

	/* Check first in this direction */
	for(i=0; i<table->len; i++){
		candidate += i;
		if(candidate->src_ip.s_addr == observed.src_ip.s_addr){
			if(candidate->dst_ip.s_addr == observed.dst_ip.s_addr){
				if(candidate->src_port == observed.src_port){
					if(candidate->dst_port == observed.dst_port){
						return i+1;
					}
				}
			}
		}
	}


	candidate = table->table;
	/* Then in the other direction */
	for(i=0; i<table->len; i++){
		candidate += i;
		if(candidate->src_ip.s_addr == observed.dst_ip.s_addr){
			if(candidate->dst_ip.s_addr == observed.src_ip.s_addr){
				if(candidate->src_port == observed.dst_port){
					if(candidate->dst_port == observed.src_port){
						return i+1;
					}
				}
			}
		}
	}

	return 0;
}

flow *get_flow(int index){
	if(index < table->len){
		return table->table+index;
	} else {
		return NULL;
	}
}

/* Adds a packet the flow's packet chain */
int add_packet(flow *f, uint8_t *p){
	const struct ip_header *ip_hdr;
	const struct tcp_header *tcp_hdr;

	packet *new_packet = malloc(sizeof(packet));

	p += ETHER_HEADER_LEN; //skip ethernet header
	ip_hdr = (struct ip_header*) p;
	int size_ip = IP_HEADER_LEN(ip_hdr);

	if (ip_hdr->proto != IPPROTO_TCP){
		return 0;
	}

	p += size_ip;	//skip IP header

	tcp_hdr = (struct tcp_header*) p;
	int size_tcp = TCP_HEADER_LEN(tcp_hdr);
	p += size_tcp;

	new_packet->seq_num = htonl(tcp_hdr->sequence_num);
	new_packet->len = htons(ip_hdr->len) - (size_ip + size_tcp);
	new_packet->data = p;
	new_packet->data_len = htons(ip_hdr->len) - (size_ip + size_tcp);
	new_packet->next = NULL;
	new_packet->incoming = 
		(ip_hdr->src.s_addr == f->src_ip.s_addr) ? 0 : 1;

	/* Find appropriate place in chain */
	if(new_packet->data_len > 0){
		packet *previous = NULL;
		packet *next = f->packet_chain;
		while(next != NULL && (next->seq_num <= new_packet->seq_num)){
			//printf("next: %u <= new: %u\n", next->seq_num, new_packet->seq_num);
			previous = next;
			next = next->next;
		}

		//place packet after current
		if(previous == NULL){
			//goes at the beginning of chain
			new_packet->next = f->packet_chain;
			f->packet_chain = new_packet;
		} else {
			new_packet->next = next;
			previous->next = new_packet;
		}
		printf("Flow: %d > %d (%s)\n", ip_hdr->src.s_addr, ip_hdr->dst.s_addr, (new_packet->incoming)? "incoming":"outgoing");
		printf("ID number: %u\n", htonl(ip_hdr->id));
		printf("Sequence number: %u\n", htonl(tcp_hdr->sequence_num));
		printf("Acknowledgement number: %u\n", htonl(tcp_hdr->ack_num));
		printf("Length: %d\n", new_packet->data_len);

	}
	
	return 0;

}

/** UTILITY **/

int update_finish_hash(flow *f, uint8_t *hs){
	EVP_MD_CTX *ctx = f->finish_md_ctx;
	//find handshake length
	const struct handshake_header *hs_hdr;
	uint8_t *p = hs;
	hs_hdr = (struct handshake_header*) p;
	uint32_t hs_len = HANDSHAKE_MESSAGE_LEN(hs_hdr);
	
	/* Now compute SHA384(secret, A+seed)
	 * can't sign update, don't have master secret yet*/
	EVP_DigestUpdate(&ctx, hs, hs_len);

	return 1;
	
}

int extract_parameters(flow *f, uint8_t *hs){
	DH *dh;
	uint8_t *p;
	long i;

	p = hs + HANDSHAKE_HEADER_LEN;

	if((dh = DH_new()) == NULL){
		return 1;
	}

	/* Extract prime modulus */
	n2s(p,i);

	if(!(dh->p = BN_bin2bn(p,i,NULL))){
		return 1;
	}
	p += i;

	/* Extract generator */
	n2s(p,i);

	if(!(dh->g = BN_bin2bn(p,i,NULL))){
		return 1;
	}
	p += i;

	/* Extract server public value */
	n2s(p,i);

	if(!(dh->pub_key = BN_bin2bn(p,i,NULL))){
		return 1;
	}

	f->dh = dh;
	printf("Param extraction: success!\n");
	return 0;
}

/* Decrypt the TLS FINISHED message
 * Verify that the data is:
 * 	PRF(master_secret, finished_label, Hash(handshake_messages))*/
int decrypt_fin(flow *f, uint8_t *hs, int32_t len, int32_t incoming){
	//const struct handshake_header *hs_hdr;
	uint8_t *p = hs;
	//EVP_MD_CTX ctx;
	//uint8_t hash[EVP_MAX_MD_SIZE];
	//int32_t hash_len;

	//EVP_MD_CTX_init(&ctx);

	//decrypt
	EVP_CIPHER_CTX *ds = (incoming) ? f->read_ctx : f->write_ctx;
	if(ds == NULL){
		printf("FAIL\n");
		return 1;
	}
	if(ds->iv[EVP_GCM_TLS_FIXED_IV_LEN] == 0){
		//fill in rest of iv
		for(int i = EVP_GCM_TLS_FIXED_IV_LEN; i< ds->cipher->iv_len; i++){
			ds->iv[i] = p[i- EVP_GCM_TLS_FIXED_IV_LEN];
		}
	}

//#ifdef KSSL_DEBUG
	{    
		fprintf(stderr, "\t\tIV: ");
		for (int i = 0; i < ds->cipher->iv_len; i++) 
			fprintf(stderr, "%02X", ds->iv[i]);
		fprintf(stderr, "\n");
	}    
//#endif                          /* KSSL_DEBUG */


	int32_t bs = EVP_CIPHER_block_size(ds->cipher);
	//padding stuff? TODO: understand this
	uint8_t buf[13];
	memset(buf, 0, 8);
	buf[8] = 0x16;
	buf[9] = 0x03;
	buf[10] = 0x03;
	buf[11] = 0x00; //len >> 8;
	buf[12] = 0x28; //len *0xff;
	printf("buf: \n");
	for(int z=0; z< 13; z++){
		printf("%02x ", buf[z]);
	}
	printf("\n");
	int32_t pad = EVP_CIPHER_CTX_ctrl(ds, EVP_CTRL_AEAD_TLS1_AAD,
			13, buf);
	printf("pad: %d\n", pad);

	printf("Decrypting (%d bytes):\n", len);
	for(int i=0; i<len; i++){
		printf("%02x ", p[i]);
	}
	printf("\n");

	int32_t n = EVP_Cipher(ds, p, p, len); //decrypt in place

	p += EVP_GCM_TLS_EXPLICIT_IV_LEN;
	len -= EVP_GCM_TLS_EXPLICIT_IV_LEN;
	//print out decrypted record
	printf("Record (decrypted %d/%d bytes): \n", n, len);
	for(int i=0; i<n; i++){
		printf("%02x ", p[i]);
	}
	printf("\n");

	/*get context
	
	//get header length
	hs_hdr = (struct handshake_header*) p;
	uint32_t fin_length = HANDSHAKE_MESSAGE_LEN(hs_hdr);
	p += HANDSHAKE_HEADER_LEN;

	//finalize hash of handshake msgs
	EVP_MD_CTX_copy_ex(&ctx, f->finish_md_ctx);
	EVP_DigestFinal(&ctx, hash, &hash_len);

	//now use pseudorandom function
	uint8_t *output = calloc(1, fin_length);
	PRF(f->master_secret, SSL3_MASTER_SECRET_SIZE, finished_label, finished_label_len, hash, hash_len, NULL, 0, NULL, 0, output, fin_length);

	//now compare
	
*/

	return 1;
}

int compute_master_secret(flow *f){
	DH *dh_srvr = NULL;
	DH *dh_clnt = NULL;
	BN_CTX *ctx;
	BN_MONT_CTX *mont = NULL;
	BIGNUM *pub_key = NULL, *priv_key = NULL;

	ctx = BN_CTX_new();

	dh_srvr = f->dh;
	dh_clnt = DHparams_dup(dh_srvr);

    uint32_t l = dh_clnt->length ? dh_clnt->length : BN_num_bits(dh_clnt->p) - 1;
    int32_t bytes = (l+7) / 8;

    uint8_t *buf = (uint8_t *)OPENSSL_malloc(bytes);
    if (buf == NULL){
        BNerr(BN_F_BNRAND, ERR_R_MALLOC_FAILURE);
        return 1;
    }

	pub_key = BN_new();
	priv_key = BN_new();

	for(int i=0; i<bytes; i++){
		buf[i] = f->key[i%16];
	}

    if (!BN_bin2bn(buf, bytes, priv_key))
		return 1;

    {
        BIGNUM *prk;

		prk = priv_key;

        if (!dh_clnt->meth->bn_mod_exp(dh_clnt, pub_key, dh_clnt->g, prk, dh_clnt->p, ctx, mont)){
			printf("FAIL\n");
            return 1;
		}
		printf("here\n");
    }

    dh_clnt->pub_key = pub_key;
    dh_clnt->priv_key = priv_key;

	// Compute master key 
	uint8_t *pre_master_secret = calloc(1, 256);//TODO: find right length

	DH_compute_key(pre_master_secret, dh_srvr->pub_key, dh_clnt);
	
	PRF(pre_master_secret, PRE_MASTER_LEN, TLS_MD_MASTER_SECRET_CONST, TLS_MD_MASTER_SECRET_CONST_SIZE, f->client_random, SSL3_RANDOM_SIZE, f->server_random, SSL3_RANDOM_SIZE, NULL, 0, f->master_secret, SSL3_MASTER_SECRET_SIZE);

	//remove pre_master_secret from memory
	memset(pre_master_secret, 0, PRE_MASTER_LEN);

	printf("master secret:\n");
	for(int i=0; i< 48; i++){
		printf("%02x ", f->master_secret[i]);
	}
	printf("\n");
	
	free(pre_master_secret);
	DH_free(dh_srvr);
	DH_free(dh_clnt);
	return 0;
}

void extract_server_random(flow *f, uint8_t *hs){

	uint8_t *p;

	p = hs + HANDSHAKE_HEADER_LEN;

	p+=2; //skip version

	memcpy(f->server_random, p, SSL3_RANDOM_SIZE);
	printf("got server random\n");

}

/* PRF using sha384, as defined in RFC 5246 */
int PRF(uint8_t *secret, int32_t secret_len,
		uint8_t *seed1, int32_t seed1_len,
		uint8_t *seed2, int32_t seed2_len,
		uint8_t *seed3, int32_t seed3_len,
		uint8_t *seed4, int32_t seed4_len,
		uint8_t *output, int32_t output_len){

	EVP_MD_CTX ctx, ctx_tmp, ctx_init;
	EVP_PKEY *mac_key;
	const EVP_MD *md = EVP_sha384();

	uint8_t A[EVP_MAX_MD_SIZE];
	size_t len, A_len;
	int chunk = EVP_MD_size(md);
	int remaining = output_len;

	uint8_t *out = output;

	EVP_MD_CTX_init(&ctx);
	EVP_MD_CTX_init(&ctx_tmp);
	EVP_MD_CTX_init(&ctx_init);
	EVP_MD_CTX_set_flags(&ctx_init, EVP_MD_CTX_FLAG_NON_FIPS_ALLOW);

	mac_key = EVP_PKEY_new_mac_key(EVP_PKEY_HMAC, NULL, secret, secret_len);

	/* Calculate first A value */
	EVP_DigestSignInit(&ctx_init, NULL, md, NULL, mac_key);
	EVP_MD_CTX_copy_ex(&ctx, &ctx_init);
	if(seed1 != NULL && seed1_len > 0){
		EVP_DigestSignUpdate(&ctx, seed1, seed1_len);
	}
	if(seed2 != NULL && seed2_len > 0){
		EVP_DigestSignUpdate(&ctx, seed2, seed2_len);
	}
	if(seed3 != NULL && seed3_len > 0){
		EVP_DigestSignUpdate(&ctx, seed3, seed3_len);
	}
	if(seed4 != NULL && seed4_len > 0){
		EVP_DigestSignUpdate(&ctx, seed4, seed4_len);
	}
	EVP_DigestSignFinal(&ctx, A, &A_len);

	//iterate until desired length is achieved
	while(remaining > 0){
		/* Now compute SHA384(secret, A+seed) */
		EVP_MD_CTX_copy_ex(&ctx, &ctx_init);
		EVP_DigestSignUpdate(&ctx, A, A_len);
		EVP_MD_CTX_copy_ex(&ctx_tmp, &ctx);
		if(seed1 != NULL && seed1_len > 0){
			EVP_DigestSignUpdate(&ctx, seed1, seed1_len);
		}
		if(seed2 != NULL && seed2_len > 0){
			EVP_DigestSignUpdate(&ctx, seed2, seed2_len);
		}
		if(seed3 != NULL && seed3_len > 0){
			EVP_DigestSignUpdate(&ctx, seed3, seed3_len);
		}
		if(seed4 != NULL && seed4_len > 0){
			EVP_DigestSignUpdate(&ctx, seed4, seed4_len);
		}
		
		if(remaining > chunk){
			EVP_DigestSignFinal(&ctx, out, &len);
			out += len;
			remaining -= len;

			/* Next A value */
			EVP_DigestSignFinal(&ctx_tmp, A, &A_len);
		} else {
			EVP_DigestSignFinal(&ctx, A, &A_len);
			memcpy(out, A, remaining);
			remaining -= remaining;
		}
	}
	return 1;
}

/* After receiving change cipher spec, calculate keys from master secret */
int init_ciphers(flow *f){

	EVP_CIPHER_CTX *r_ctx;
	EVP_CIPHER_CTX *w_ctx;
	const EVP_CIPHER *c = EVP_aes_256_gcm();

	/* Generate Keys */
	uint8_t *write_mac, *write_key, *write_iv;
	uint8_t *read_mac, *read_key, *read_iv;
	int32_t mac_len, key_len, iv_len;

	key_len = EVP_CIPHER_key_length(c);
	iv_len = EVP_CIPHER_iv_length(c); //EVP_GCM_TLS_FIXED_IV_LEN;
	mac_len = EVP_MD_size(EVP_get_digestbyname(SN_sha384));
	int32_t total_len = key_len + iv_len + mac_len;
	total_len *= 2;
	uint8_t *key_block = calloc(1, total_len);

	PRF(f->master_secret, SSL3_MASTER_SECRET_SIZE,
			TLS_MD_KEY_EXPANSION_CONST, TLS_MD_KEY_EXPANSION_CONST_SIZE,
			f->server_random, SSL3_RANDOM_SIZE,
			f->client_random, SSL3_RANDOM_SIZE,
			NULL, 0,
			key_block, total_len);

	printf("keyblock:\n");
	for(int i=0; i< total_len; i++){
		printf("%02x ", key_block[i]);
	}
	printf("\n");

	iv_len = EVP_GCM_TLS_FIXED_IV_LEN;
	
	write_key = key_block;
	read_key = key_block + key_len;
	write_iv = key_block + 2*key_len;
	read_iv = key_block + 2*key_len + iv_len;
	write_mac = key_block + 2*key_len + 2*iv_len;
	read_mac = key_block + 2*key_len + 2*iv_len + mac_len;

	/* Initialize Cipher Contexts */
	r_ctx = EVP_CIPHER_CTX_new();
	w_ctx = EVP_CIPHER_CTX_new();
	EVP_CIPHER_CTX_init(r_ctx);
	EVP_CIPHER_CTX_init(w_ctx);

//#ifdef KSSL_DEBUG
    {
        int i;
        fprintf(stderr, "EVP_CipherInit_ex(r_ctx,c,key=,iv=,which)\n");
        fprintf(stderr, "\tkey= ");
        for (i = 0; i < c->key_len; i++)
            fprintf(stderr, "%02x", read_key[i]);
        fprintf(stderr, "\n");
        fprintf(stderr, "\t iv= ");
        for (i = 0; i < c->iv_len; i++)
            fprintf(stderr, "%02x", read_iv[i]);
        fprintf(stderr, "\n");
    }
//#endif                          /* KSSL_DEBUG_*/
    {
        int i;
        fprintf(stderr, "EVP_CipherInit_ex(w_ctx,c,key=,iv=,which)\n");
        fprintf(stderr, "\tkey= ");
        for (i = 0; i < c->key_len; i++)
            fprintf(stderr, "%02x", write_key[i]);
        fprintf(stderr, "\n");
        fprintf(stderr, "\t iv= ");
        for (i = 0; i < c->iv_len; i++)
            fprintf(stderr, "%02x", write_iv[i]);
        fprintf(stderr, "\n");
    }
//#endif                          /* KSSL_DEBUG */


	EVP_CipherInit_ex(r_ctx, c, NULL, read_key, NULL, 0);
	EVP_CipherInit_ex(w_ctx, c, NULL, write_key, NULL, 0);
	EVP_CIPHER_CTX_ctrl(r_ctx, EVP_CTRL_GCM_SET_IV_FIXED, EVP_GCM_TLS_FIXED_IV_LEN, read_iv);
	EVP_CIPHER_CTX_ctrl(w_ctx, EVP_CTRL_GCM_SET_IV_FIXED, EVP_GCM_TLS_FIXED_IV_LEN, write_iv);
	fprintf(stderr, "\t\tIV: ");
	for (int i = 0; i < r_ctx->cipher->iv_len; i++) 
		fprintf(stderr, "%02X", r_ctx->iv[i]);
	fprintf(stderr, "\n");
		fprintf(stderr, "\t\tIV: ");
		for (int i = 0; i < w_ctx->cipher->iv_len; i++) 
			fprintf(stderr, "%02X", w_ctx->iv[i]);
		fprintf(stderr, "\n");
	f->read_ctx = r_ctx;
	f->write_ctx = w_ctx;
}