#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <regex.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <netdb.h>
#include <unistd.h>
#include <pthread.h>
#include "relay.h"
#include "slitheen.h"
#include "flow.h"
#include "crypto.h"

/** Called when a TLS application record is received for a
 *  tagged flow. Upstream packets will be checked for covert
 *  requests to censored sites, downstream packets will be
 *  replaced with data from the censored queue or with garbage
 *
 *  Inputs:
 *  	f: the tagged flow
 *  	info: the processed received application packet
 *
 *  Output:
 *  	0 on success, 1 on failure
 */
int replace_packet(flow *f, struct packet_info *info){

	if (info->tcp_hdr == NULL){
		return 0;
	}

#ifdef DEBUG
	fprintf(stderr,"Flow: %d > %d (%s)\n", info->ip_hdr->src.s_addr, info->ip_hdr->dst.s_addr, (info->ip_hdr->src.s_addr != f->src_ip.s_addr)? "incoming":"outgoing");
	fprintf(stderr,"ID number: %u\n", htonl(info->ip_hdr->id));
	fprintf(stderr,"Sequence number: %u\n", htonl(info->tcp_hdr->sequence_num));
	fprintf(stderr,"Acknowledgement number: %u\n", htonl(info->tcp_hdr->ack_num));
#endif

	if(info->app_data_len <= 0){
		return 0;
	}

	/* if outgoing, decrypt and look at header */
	if(info->ip_hdr->src.s_addr == f->src_ip.s_addr){
		read_header(f, info);
		return 0;
	} else {

#ifdef DEBUG
		printf("Current sequence number: %d\n", f->seq_num);
		printf("Received sequence number: %d\n", htonl(tcp_hdr->sequence_num));
#endif

		uint32_t offset = htonl(info->tcp_hdr->sequence_num) - f->seq_num;
		if(offset == 0)
			f->seq_num += info->app_data_len;
		/* if incoming, replace with data from queue */
		//if(htonl(tcp_hdr->sequence_num) >= f->seq_num){
			replace_contents(f, offset, info);
		//}//TODO: need to do something about replaying packets (maybe store previously sent data??


#ifdef DEBUG //TODO: fix
		uint8_t *p = (uint8_t *) info->tcp_hdr;
		fprintf(stdout, "ip hdr length: %d\n", htons(info->ip_hdr->len));
		fprintf(stdout, "Injecting the following packet:\n");
		for(int i=0; i< htons(info->ip_hdr->len); i++){
			fprintf(stdout, "%02x ", p[i]);
		}
		fprintf(stdout, "\n");
		fflush(stdout);
#endif

	}
	return 0;

}

/** Reads the HTTP header of upstream data and searches for
 *  a covert request in the x-ignore header. Sends this
 *  request to the indicated site and saves the response to
 *  the censored queue
 *
 *  Inputs:
 *  	f: the tagged flow
 *  	info: the processed received packet
 *
 *  Ouput:
 *  	0 on success, 1 on failure
 */
int read_header(flow *f, struct packet_info *info){
	uint8_t *p = info->app_data;

	if (info->tcp_hdr == NULL){
		return 0;
	}

	struct record_header *record_hdr = (struct record_header*) p;
	uint32_t record_length = RECORD_LEN(record_hdr);

	uint8_t *decrypted_data = calloc(1, info->app_data_len);

	p+= RECORD_HEADER_LEN;

	memcpy(decrypted_data, p, record_length);

	if(!encrypt(f, decrypted_data, decrypted_data, record_length, 0, record_hdr->type, 0)){
		fprintf(stdout,"decryption failed\n");
		return 0;
	}

	if(record_hdr->type == 0x15){
		printf("received alert\n");
		for(int i=0; i<record_length; i++){
			printf("%02x ", decrypted_data[i]);
		}
		fflush(stdout);
	}

	//TODO: re-write this to take a SOCKS connection request

	/* search through decrypted data for x-ignore */
	regex_t r;
	const char *regex_text = "x-ignore";
	const char *match = (char *) decrypted_data;
	if(regcomp(&r, regex_text, REG_EXTENDED|REG_NEWLINE)){
		printf("could not compile regex\n");
		return 0;
	}

	regmatch_t m;
	if(regexec(&r, match, 1, &m, 0)){
		printf("Upstream data: no x-ignore header\n");
		return 0;
	} 
	uint8_t *message = decrypted_data;
	//remove escape characters
	for(int i=m.rm_eo+2; i< strlen(match); i++){
		if(match[i] != '\\'){
			*(message++) = match[i];
		} else if (match[i+1] == 'r') {
			*(message++) = '\r';
			i++;
		} else if (match[i+1] == 'n') {
			*(message++) = '\n';
			i++;
		}
	}
	*message = '\0';
	message = decrypted_data;

	regex_t r2;
	const char *regex_text2 = "host";
	const char *match2 = (char *) decrypted_data;
	regmatch_t m2;
	if(regcomp(&r2, regex_text2, REG_EXTENDED|REG_NEWLINE)){
		printf("could not compile regex\n");
		return 0;
	}
	if(regexec(&r2, match2, 1, &m2, 0)){
		printf("no host found\n");
		return 0;
	}
	uint8_t server[50];
	for(int i=m2.rm_eo+2; i< strlen(match2); i++){
		if(match2[i] != '\n'){
			server[i-m2.rm_eo-2] = match2[i];
		} else {
			server[i-m2.rm_eo-2] = '\0';
			break;
		}
	}

	printf("Collecting data from censored site: %s\n", server);
	//If a thread for this stream id exists, get the thread info
		//TODO: fill this in

	//Else, spawn a thread to handle the proxy to this site
	pthread_t *proxy_thread = calloc(1, sizeof(pthread_t));
	int32_t pipefd[2];
	if(pipe(pipefd) < 0){
		printf("Failed to create pipe for new thread\n");
		return 1;
	}

	struct proxy_thread_data *thread_data = 
		calloc(1, sizeof(struct proxy_thread_data));
	thread_data->f = f;
	memcpy(thread_data->server, server, 50);
	printf("Collecting data from censored site: %s\n", thread_data->server);
	thread_data->pipefd = pipefd[0];
	
	pthread_create(proxy_thread, NULL, proxy_covert_site, (void *) thread_data);

	//save a reference to the proxy threads in the flow
	

	//Now pipe information to the selected or created thread
	//TODO: message may not be a string. fix.
	int32_t bytes_written = write(pipefd[1], message,
			strlen( (const char *) message));
	if(bytes_written < strlen( (const char *) message)){
		printf("failed to write all bytes to pipe\n");
	}

	return 0;

}

/** Called by spawned pthreads in read_header to send upstream
 *  data to the censored site and receive responses. Downstream
 *  data is stored in the flow's censored_queue. Function and
 *  thread will terminate when the client closes the connection
 *  to the covert destination
 *
 *  Input:
 *  	A struct that contains the following information:
 *  	- the tagged flow
 *  	- the name of the server
 *  	- the read end of the pipe
 *
 */
void *proxy_covert_site(void *data){

	struct proxy_thread_data *thread_data =
		(struct proxy_thread_data *) data;

	flow *f = thread_data->f;

	int32_t buffer_len = BUFSIZ;
	uint8_t *buffer = calloc(1, BUFSIZ);
	/* Send GET request to site */
	printf("Collecting data from censored site: %s\n", thread_data->server);

	struct hostent *host;
	host = gethostbyname((const char *) thread_data->server);

	if(host == NULL){
		printf("gethostbyname failed\n");
		return 0;
	}

	int32_t handle = socket(AF_INET, SOCK_STREAM, 0);
	if(handle < 0){
		printf("error: constructing socket failed\n");
		return 0;
	}
	
	struct sockaddr_in dest;
	dest.sin_family = AF_INET;
	dest.sin_port = htons(80);
	dest.sin_addr = *((struct in_addr *) host->h_addr);
	bzero (&(dest.sin_zero), 8);
	
	int32_t error = connect (handle, (struct sockaddr *) &dest, sizeof (struct sockaddr));
	
	if(error <0){
		printf("error connecting\n");
		return 0;
	}

	//now select on reading from the pipe and from the socket
	for(;;){
		fd_set readfds;
		fd_set writefds;

		int32_t nfds = (handle > thread_data->pipefd) ?
			handle +1 : thread_data->pipefd + 1;

		FD_ZERO(&readfds);
		FD_ZERO(&writefds);

		FD_SET(thread_data->pipefd, &readfds);
		FD_SET(handle, &readfds);
		FD_SET(handle, &writefds);

		if (select(nfds, &readfds, &writefds, NULL, NULL) < 0){
			printf("select error\n");
		}

		if(FD_ISSET(thread_data->pipefd, &readfds) && FD_ISSET(handle, &writefds)){
			//we have upstream data ready for writing
			printf("Passing along upstream data\n");

			int32_t bytes_read = read(thread_data->pipefd, buffer, buffer_len);
			buffer[buffer_len] = '\0';//TODO: remove w/ print
			printf("Read from pipe:\n %s\n", buffer);

			if(bytes_read > 0){
				int32_t bytes_sent = send(handle, buffer,
						bytes_read, 0);
				if( bytes_sent < 0){
					printf("error sending request\n");
					break;
				} else if (bytes_sent < bytes_read){
					//TODO: should update buffer and keep
					//track of length of upstream data
					printf("sent less than full upstream bytes\n");
					break;
				}
			}

		}
		
		if (FD_ISSET(handle, &readfds)){
			//we have downstream data read for saving
			int32_t bytes_read;
			uint8_t *buf = calloc(1, BUFSIZ);
			bytes_read = recv(handle, buf, BUFSIZ, 0);
			if(bytes_read <= 0){
				break;
			}
			if(bytes_read > 0){

				//make a new queue block
				queue_block *new_block = calloc(1, sizeof(queue_block));
				new_block->len = bytes_read;
				new_block->offset = 0;
				new_block->data = buf;
				new_block->next = NULL;
				if(f->censored_queue == NULL)
					f->censored_queue = new_block;
				else{
					queue_block *last = f->censored_queue;
					while(last->next != NULL)
						last = last->next;
					last->next = new_block;
				}
			} else {
				printf("read 0 bytes\n");
			}

		}
	}
	free(thread_data);
	free(buffer);
	close(handle);
	return 0;
}

/** Replaces downstream record contents with data from the
 *  censored queue, padding with garbage bytes if no more
 *  censored data exists.
 *
 *  Inputs: 
 *  	f: the tagged flow
 *  	data: a pointer to the received packet's application
 *  		data
 *  	data_len: the length of the	packet's application data
 *  	offset: if the packet is misordered, the number of
 *  		application-level bytes in missing packets
 *
 *  Output:
 *  	Returns 0 on sucess 
 */
int replace_contents(flow *f, int32_t offset, struct packet_info *info){

	printf("Replacing contents of downstream data\n");

	uint8_t *p = info->app_data;

	int32_t tmp_len = info->app_data_len;

	//step 1: replace record contents
	//note: encrypted message will be original message size + EVP_GCM_TLS_EXPLICIT_IV_LEN + 16 byte pad
	//first check to see if there's anything in the outbox
	if(f->outbox_len > 0){

#ifdef DEBUG
		if(f->outbox_len < info->app_data_len){
			printf("Next record:\n");
			for(int i=0; i< RECORD_HEADER_LEN; i++){
				printf("%02x ", p[f->outbox_len+i]);
			}
			printf("\n");
		} else {
			printf("Outbox takes up entire packet\n");
		}
#endif

		if(tmp_len >= f->outbox_len){
			memcpy(p, f->outbox, f->outbox_len);

			p += f->outbox_len;
			tmp_len -= f->outbox_len;
			f->outbox_len = 0;
			free(f->outbox);
		} else {
			memcpy(p, f->outbox, tmp_len);
			uint8_t *tmp = calloc(1, f->outbox_len - tmp_len);
			f->outbox_len -= tmp_len;
			memcpy(tmp, f->outbox + tmp_len, f->outbox_len);
			free(f->outbox);
			f->outbox = tmp;
			tmp_len -= tmp_len;
		}
	}

	while(tmp_len > 0){

		struct record_header *record_hdr = (struct record_header*) p;
		uint32_t record_length = RECORD_LEN(record_hdr);

#ifdef DEBUG
		fprintf(stdout, "Record:\n");
		for(int i=0; i< RECORD_HEADER_LEN; i++){
			printf("%02x ", p[i]);
		}
		printf("\n");
#endif

		p += RECORD_HEADER_LEN;
		if(record_hdr->type != 0x17){
			//TODO: might need to decrypt and re-encrypt
			//printf("received non-application data\n");
			tmp_len -= (record_length+ RECORD_HEADER_LEN);
			p += record_length;
			continue;
		}

		uint8_t *new_record = calloc(1, record_length);
		memcpy(new_record, p, record_length);
		uint8_t *tmp_p = new_record;

		tmp_p += EVP_GCM_TLS_EXPLICIT_IV_LEN;
		
		struct slitheen_header *sl_hdr = (struct slitheen_header *) tmp_p;
		sl_hdr->marker = 0x01;
		sl_hdr->version = 0x01;
		sl_hdr->len = 0x00;
		int32_t remaining = record_length - (SLITHEEN_HEADER_LEN
				+ EVP_GCM_TLS_EXPLICIT_IV_LEN + 16);
		tmp_p += SLITHEEN_HEADER_LEN;
		//Fill as much as we can from the censored_queue
		while((remaining > 0) && f->censored_queue != NULL){
			int32_t block_length = f->censored_queue->len;
			int32_t offset = f->censored_queue->offset;

#ifdef DEBUG
			printf("Censored queue is at %p.\n", f->censored_queue);
			printf("This block as %d bytes left\n", block_length - offset);
			printf("We need %d bytes\n", remaining);
#endif
			
			if(block_length > offset + remaining){
				//use part of the block, update offset
				memcpy(tmp_p, f->censored_queue->data+offset, remaining);
				f->censored_queue->offset += remaining;
				tmp_p += remaining;
				sl_hdr->len += remaining;
				remaining -= remaining;
			} else {
				//use all of the block and free it
				memcpy(tmp_p, f->censored_queue->data+offset, block_length - offset);

				free(f->censored_queue->data);
				f->censored_queue = f->censored_queue->next;

				tmp_p += (block_length - offset);
				sl_hdr->len += (block_length - offset);
				remaining -= (block_length - offset);
			}
		}
		sl_hdr->len = htons(sl_hdr->len);
		//now, if we need more data, fill with garbage
		if(remaining >0 ){
			//TODO: note, we may also be receiving misordered packets. Take Ian's suggestion into account here
			memset(tmp_p, 'A', remaining);
		}

		tmp_p = new_record;

#ifdef DEBUG
		fprintf(stdout, "copied %d data and %d garbage bytes\n", ntohs(sl_hdr->len), remaining);
		printf("Slitheen header\n");
		for(int i=0; i<4; i++)
			printf("%02x ", tmp_p[EVP_GCM_TLS_EXPLICIT_IV_LEN+i]);
		printf("\n");
#endif

		//step 3: encrypt new record
		int32_t success;
		if((success = encrypt(f, tmp_p, tmp_p, record_length-16, 1, 0x17, 1))< 0){
			fprintf(stdout,"encryption failed\n");
			return 0;
		}

		//copy new record into packet
		if(record_length +RECORD_HEADER_LEN > tmp_len){
			//We have a partial record
			memcpy(p, new_record, tmp_len - RECORD_HEADER_LEN);
			f->outbox_len = record_length - (tmp_len - RECORD_HEADER_LEN);
			//save left-overs in outbox
			f->outbox = calloc(1, f->outbox_len);
			memcpy(f->outbox, new_record + (tmp_len - RECORD_HEADER_LEN),
					f->outbox_len);
			free(new_record);
		} else {
			memcpy(p, new_record, record_length);
			free(new_record);
		}

#ifdef DEBUG
		//check to see if next record still exists
		if(info->app_data_len > record_length + RECORD_HEADER_LEN){
			printf("Next record:\n");
			for(int i=0; i< RECORD_HEADER_LEN; i++){
				printf("%02x ", p[record_length+i]);
			}
			printf("\n");
		} else {
			printf("No extra record: %d <= %d + %d\n", data_len, record_length, RECORD_HEADER_LEN);
		}
#endif

		tmp_len -= record_length+ RECORD_HEADER_LEN;
		p += record_length;

	}
	//step 4: recompute TCP checksum
	tcp_checksum(info);

	return 0;
}

/** Computes the TCP checksum of the data according to RFC 793
 *  sum all 16-bit words in the segment, padd the last word if
 *  needed
 *
 *  there is a pseudo-header prefixed to the segment and
 *  included in the checksum:
 *
 *         +--------+--------+--------+--------+
 *         |           Source Address          |
 *         +--------+--------+--------+--------+
 *         |         Destination Address       |
 *         +--------+--------+--------+--------+
 *         |  zero  |  PTCL  |    TCP Length   |
 *         +--------+--------+--------+--------+
 */
uint16_t tcp_checksum(struct packet_info *info){

	uint16_t tcp_length = info->app_data_len + info->size_tcp_hdr;
	struct in_addr src = info->ip_hdr->src;
	struct in_addr dst = info->ip_hdr->dst;
	uint8_t proto = IPPROTO_TCP;

	//set the checksum to zero
	info->tcp_hdr->chksum = 0;
	
	//sum pseudoheader
	uint32_t sum = (ntohl(src.s_addr)) >> 16;
	sum += (ntohl(src.s_addr)) &0xFFFF;
	sum += (ntohl(dst.s_addr)) >> 16;
	sum += (ntohl(dst.s_addr)) & 0xFFFF;
	sum += proto;
	sum += tcp_length;

	//sum tcp header (with zero-d checksum)
	uint8_t *p = (uint8_t *) info->tcp_hdr;
	for(int i=0; i < info->size_tcp_hdr; i+=2){
		sum += (uint16_t) ((p[i] << 8) + p[i+1]);
	}

	//now sum the application data
	p = info->app_data;
	for(int i=0; i< info->app_data_len-1; i+=2){
		sum += (uint16_t) ((p[i] << 8) + p[i+1]);
	}
	if(info->app_data_len %2 != 0){
		sum += (uint16_t) (p[info->app_data_len - 1]) << 8;
	}
	//now add most significant to last significant bits
	sum = (sum >> 16) + (sum & 0xFFFF);
	//now subtract from 0xFF
	sum = 0xFFFF - sum;

	//set chksum to calculated value
	info->tcp_hdr->chksum = ntohs(sum);
	return (uint16_t) sum;
}