#include <pcap.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <pthread.h>
#include "ptwist.h"
#include "rserv.h"
#include "flow.h"

#define macaddr "08:00:27:e8:9d:d4"

/* Ethernet addresses are 6 bytes */
#define ETHER_ADDR_LEN	6
#define ETHER_HEADER_LEN  2*ETHER_ADDR_LEN + 2

/* Definitions for parsing packet data */
struct ip_header {
	u_char versionihl;		/* Version >> 4 | IHL & 0x0f */
	u_char dscpecn;		/* DSCP >> 2 | ECN  & 0x03 */
	u_short len;		/* Total Length */
	u_short id;		/* Identification */
	u_short flagsoff;		/* Flags >> 13 | Fragment Offset & 0x1fff */
#define RF 0x8000		/* Reserved; must be zero */
#define DF 0x4000		/* Dont Fragment */
#define MF 0x2000		/* More Fragments */
	u_char ttl;		/* Time To Live */
	u_char proto;		/* Protocol */
	u_short chksum;		/* Header Checksum */
	struct in_addr src, dst; /* Source and Destination addresses */
};
#define IP_HEADER_LEN(ip)		(((ip)->versionihl) & 0x0f)*4

struct tcp_header {
	u_short src_port;	/* source port */
	u_short dst_port;	/* destination port */
	u_int sequence_num;		/* sequence number */
	u_int ack_num;		/* acknowledgement number */
	u_char offset_res_ns;	/*Data offset >> 4 |  res >> 1 | NS 0x01 */
	u_char flags;			/* Flags */
#define FIN 0x01
#define RST 0x04
	u_short win_size;		/* Window size*/
	u_short chksum;		/* Checksum */
	u_short urg;		/* Urgent pointer */
};
#define TCP_HEADER_LEN(tcp)		(((tcp)->offset_res_ns) >> 4)*4

struct tls_header {
	u_char type; /* Content Type */
#define CCS 0x14
#define A   0x15
#define HS  0x16
#define APP 0x17
#define HB  0x18
	u_short version; /* Version */
	u_short len; /* Length */
	u_char msg; /* Message Type */
#define CLIENT_HELLO 0x01
#define FINISHED 0x14
};
#define RECORD_HEADER_LEN 5
#define CLIENT_HELLO_HEADER_LEN 6

struct sniff_args {
	char *readdev;
	char *writedev;
	char *filter;
};

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);
void *sniff_packets(void *);

void check_handshake(const struct tcp_header *tcp_hdr, flow f, unsigned char *ptr){

	FILE *fp;
	int res, i, code;
	unsigned char *p;

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

	p = ptr;
	code = p[0];

	printf("handshake code: %d\n", code);

	if (code == 0x01){
		p += CLIENT_HELLO_HEADER_LEN;
		//now pointing to hello random :D
		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);
			strncpy((char *) f.key, (char *) key, 16);

			//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 */
			printf("saving flow..\n");
			add_flow(f);
			printf("saved\n");

		}
	} else {
		int index = check_flow(f);
		update_flow(index, code);
	}

}

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;
	int index;

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

	u_int size_ip;
	u_int size_tcp;

	flow newFlow;

	handle = pcap_open_live(writedev, BUFSIZ, 1, 1000, errbuf);
	if (handle == NULL){
		fprintf(stderr, "Couldn't open device %s: %s\n", writedev, errbuf);
	}
	/* check for clientHello */
	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;

		if (p[0] == 0x16){ /* This is a TLS handshake */
			p += RECORD_HEADER_LEN;
			check_handshake(tcp_hdr, newFlow, p);
		}
	}

	if((index = check_flow(newFlow))){
		flow *observed = get_flow(index-1);

		/* 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->tls_state & TLS_FINISHED){
			//decrypt packet?
			printf("TLS finished received.\n");
		} else {
			//check to see if tls_finished message
			if((observed->tls_state & TLS_NEW_SESS) && !observed->encrypted){
				//packet should be encrypted
				observed->encrypted = 1;
			} else if(observed->encrypted){ /* decrypt tls finished message */
				printf("need to decrypt this finished message\n");
				update_flow(index, 20);
			}

		}

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

	} else {

		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();

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