#include <pcap.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <pthread.h>
#include <openssl/ssl.h>
#include "rserv.h"
#include "flow.h"
#include "slitheen.h"
#include "util.h"

void got_packet(u_char *args, const struct pcap_pkthdr *header, const u_char *packet);
void check_handshake(const struct tcp_header *tcp_hdr, flow f, unsigned char *p, u_short length);
void *sniff_packets(void *);
void process_packet(const u_char *packet);

/** Checks a handshake message to see if it is tagged or a
 *  recognized flow.
 *  Inputs: The tcp header for the message, its associated flow,
 *  a pointer to the handshake header, and the length of
 *  the record.
 */
void check_handshake(const struct tcp_header *tcp_hdr, flow f, unsigned char *ptr, u_short length){

	FILE *fp;
	int res, i, code;
	//u_int size_hs;
	unsigned char *p;
	unsigned char *hello_rand;
	const struct handshake_header *handshake_hdr;

    byte privkey[PTWIST_BYTES];
	byte key[16];

	p = ptr;
	handshake_hdr = (struct handshake_header*) p;

	code = handshake_hdr->type;
	//size_hs = HANDSHAKE_MESSAGE_LEN(handshake_hdr);

	if (code == 0x01){
		p += CLIENT_HELLO_HEADER_LEN;
		//now pointing to hello random :D
		hello_rand = p;
		p += 4; //skipping time bytes
		/* Load the private key */
		fp = fopen("privkey", "rb");
		if (fp == NULL) {
			perror("fopen");
			exit(1);
		}
		res = fread(privkey, PTWIST_BYTES, 1, fp);
		if (res < 1) {
			perror("fread");
			exit(1);
		}
		fclose(fp);

		/* check tag*/ 
		res = check_tag(key, privkey, p, (const byte *)"context", 7);
		if (res) {
			printf("Untagged\n");
		} else {
			fp = fopen("tags", "wb");
			if (fp == NULL) {
				perror("fopen");
				exit(1);
			}
			//Write tag to file
			for(i=0; i< 28; i++){
				fprintf(fp, "%02x ", p[i]);
			}
			fclose(fp);

			//Write key to file
			fp = fopen("sharedkey", "wb");
			if (fp == NULL) {
				perror("fopen");
				exit(1);
			}
			for(i=0; i<16;i++){
			    fprintf(fp, "%02x", key[i]);
			}
			fclose(fp);

			/* Save flow in table */
			flow *flow_ptr = add_flow(f);
			for(int i=0; i<16; i++){
				flow_ptr->key[i] = key[i];
			}

			printf("KEY \n");
			for(int i=0; i< 16; i++){
				printf("%02x ", flow_ptr->key[i]);
			}
			printf("\n");
			memcpy(flow_ptr->client_random, hello_rand, SSL3_RANDOM_SIZE);
			
			printf("Saved new flow\n");

		}
	}

}

void got_packet(u_char *args, const struct pcap_pkthdr *header, const u_char *packet){
	pcap_t *handle;
	char errbuf[BUFSIZ];
	char *writedev = (char *) args;
	unsigned char *p;

	const struct ip_header *ip_hdr;
	//const struct tcp_header *tcp_hdr;

	handle = pcap_open_live(writedev, BUFSIZ, 1, 1000, errbuf);
	if (handle == NULL){
		fprintf(stderr, "Couldn't open device %s: %s\n", writedev, errbuf);
	}

	p = (unsigned char *) packet;
	p += ETHER_HEADER_LEN; //skip ethernet header
	ip_hdr = (struct ip_header*) p;
	p += IP_HEADER_LEN(ip_hdr);//TODO: remove
	//tcp_hdr = (struct tcp_header*) p;//TODO: remove

	/* Check to make sure it is an IP packet */
	if( (ip_hdr->versionihl >>4) != 4)
		goto end;

	/* Packet may be split over multiple frames. In this case,
	 * reconstruct the packet before trying to read TLS records
	 */
	if((htons(ip_hdr->flagsoff)&MF) || (htons(ip_hdr->flagsoff) &0x1fff) ){ /* this is a fragment */
		printf("MF: %d, OFF: %d.\n", htons(ip_hdr->flagsoff)&MF, htons(ip_hdr->flagsoff)&0x1fff);
		printf("Received packet fragment.\n");
		printf("%d packet bytes: \n", IP_HEADER_LEN(ip_hdr));
		for(int i=0; i< IP_HEADER_LEN(ip_hdr); i++){
			printf("%02x ", p[i]);
		}
		printf("\n");
		u_char *complete_packet;
		//save packet fragment
		complete_packet = malloc(65535);

		int success = report_fragment(p, complete_packet, header->len);
		if(success){
			process_packet(complete_packet);
		}
		/* TODO: handle errors */
		
	} else { /*not a fragment, add to packet chain */
		process_packet(packet);
	}
	
end:

	/* Hand packet to relay module */
	if((pcap_inject(handle, packet, header->len)) < 0 ){
		fprintf(stderr, "Error: %s\n", pcap_geterr(handle));
	}

	pcap_close(handle);

}

void usage(void){
	printf("Usage: slitheen-proxy [internal network interface] [NAT interface]\n");
}

int main(int argc, char *argv[]){
	pthread_t t1, t2;
	char filter1[33] = "ether src host 08:00:27:e8:9d:d4";
	char filter2[33] = "ether src host 08:00:27:e8:9d:d4";

	char *dev1 = NULL; /* Device that leads to the internal network */
	char *dev2 = NULL; /* Device that leads out to the world */

	struct sniff_args outbound;
	struct sniff_args inbound;

	if (argc != 3) { 
		usage();
		return(2);
	}
	dev1 = argv[1];
	dev2 = argv[2];

	snprintf(filter1, 33, "ether src host %s", macaddr);
	snprintf(filter2, 33, "ether dst host %s", macaddr);

	init_flow_table();
	init_fragment_table();

	/* Create threads */
	outbound.readdev = dev1;
	outbound.writedev = dev2;
	outbound.filter = filter1;
	inbound.readdev = dev2;
	inbound.writedev = dev1;
	inbound.filter = filter2;
	pthread_create(&t1, NULL, sniff_packets, (void *) &outbound);
	pthread_create(&t2, NULL, sniff_packets, (void *) &inbound);

	pthread_join(t1, NULL);
	pthread_join(t2, NULL);
	
	return(0);
}

void *sniff_packets(void *args){
	pcap_t *handle;
	char errbuf[BUFSIZ];
	struct bpf_program fp;
	bpf_u_int32 mask;
	bpf_u_int32 net;

	char *readdev, *writedev, *filter;
	struct sniff_args *arg_st = (struct sniff_args *) args;
	readdev = arg_st->readdev;
	writedev = arg_st->writedev;
	filter = arg_st->filter;

	if (pcap_lookupnet(readdev, &net, &mask, errbuf) == -1){
		fprintf(stderr, "Can't get netmask for device %s\n", readdev);
		exit(2);
	}

	handle = pcap_open_live(readdev, BUFSIZ, 1, 1000, errbuf);
	if (handle == NULL){
		fprintf(stderr, "Couldn't open device %s: %s\n", readdev, errbuf);
	}

	if(pcap_datalink(handle) != DLT_EN10MB) {
		fprintf(stderr, "Device %s does not provide Ethernet headers - not supported\n", readdev);
		exit(2);
	}

	if(pcap_compile(handle, &fp, filter, 0 , net) == -1){
		fprintf(stderr, "Couldn't parse filter %s: %s\n", filter, pcap_geterr(handle));
		exit(2);
	}

	if (pcap_setfilter(handle, &fp) == -1) {
		fprintf(stderr, "Couldn't install filter %s: %s\n", filter, pcap_geterr(handle));
		exit(2);
	}

	/*callback function*/
	pcap_loop(handle, -1, got_packet, (unsigned char *) writedev);

	/*Sniff a packet*/
	pcap_close(handle);

	return NULL;
}

/* This function receives a full ip packet and then:
 * 	1) identifies the flow
 * 	2) adds the packet to the flow's data chain
 * 	3) updates the packet's state
 */
void process_packet(const u_char *packet){
	unsigned char *p;
	int index;

	const struct ip_header *ip_hdr;
	const struct tcp_header *tcp_hdr;
	const struct record_header *record_hdr;

	u_int size_ip;
	u_int size_tcp;

	flow newFlow;

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

	if (ip_hdr->proto == IPPROTO_TCP){
		p += size_ip;	//skip IP header

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

		newFlow.src_ip = ip_hdr->src;
		newFlow.dst_ip = ip_hdr->dst;
		newFlow.src_port = tcp_hdr->src_port;
		newFlow.dst_port = tcp_hdr->dst_port;

		newFlow.seq_num = tcp_hdr->sequence_num;
		record_hdr = (struct record_header*) p;

		/* Checks to see if this is a possibly tagged hello msg */
		if (record_hdr->type == HS){ /* This is a TLS handshake */
			//printf("record type: %d.\n", record_hdr->type);
			//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));
			//printf("record header bytes:\n");
			//for(i=0; i<RECORD_HEADER_LEN; i++){
			//	printf("%02x ", p[i]);
			//}
			//printf("\n");
			p += RECORD_HEADER_LEN;
			check_handshake(tcp_hdr, newFlow, p, RECORD_LEN(record_hdr));
			//printf("Length of packet: %d.\n", htons(ip_hdr->len));
		}
	}

	/* Now if flow is in table, update state */
	if((index = check_flow(newFlow))){
		flow *observed = get_flow(index-1);


		/* Pass data to packet chain */
		add_packet(observed, packet);

		/* Update flow state */
		if(observed->packet_chain != NULL){
			update_flow(observed);
		}

		/* Update TCP state */
		if(tcp_hdr->flags & (FIN | RST) ){
			/* Remove flow from table, connection ended */
			remove_flow(index);
		}

		/* Check to see if TLS finished message has passed */
		if(observed->state == TLS_FINISHED){
			//decrypt packet?
			printf("TLS finished received.\n");
		} else {
			//check to see if tls_finished message
			if(observed->in_encrypted && observed->out_encrypted){ /* decrypt tls finished message */
				printf("MESSAGE ENCRYPTED\n");
			}

		}
	}

}