/* Library dependencies :
	libpcap (sudo apt-get install libpcap0.8-dev)
*/

/* Basics */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>

/* Args parsing and basename() for usage */
#include <getopt.h>
#include <libgen.h>

/* UDP stuff */
#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <unistd.h>

/* Packet capture stuff */
#include <pcap.h>
#include <sys/time.h>

/* Time management (for bwlimit) */
#include <time.h>

/* Option default values (as strings) */
#define DEFAULT_FILTER "not ( icmp[icmptype]=icmp-unreach or (udp and dst %s and port %s) )"
#define DEFAULT_HOST "127.0.0.1"
#define DEFAULT_PORT "37008"
#define DEFAULT_SNAPLEN "70"

/* MAX_TZSP_PAYLOAD is for not trying to send TZSP packets greater than the interface MTU 
   MAX_TZSP_PAYLOAD =  MTU - IP header - UDP header - TZSP header 
   I am assuming the following :
     - Layer 2 overhead is the same for the captured packet and the transmitted packet
     - MTU is assumed to be 1500 for now
*/
#define MAX_TZSP_PAYLOAD (1500-20-8-16)
#define UDP_SEND_BUFLEN 1500
#define MIN_BYTES_TO_CAPTURE 16
#define MAX_BYTES_TO_CAPTURE (1500-20-8-16)


/*
TODO List
  * Resolution DNS host -> IP (c'est l'ip qu'il faut dans le filtre et pas le host !!)
  * Implémenter une bwlimit en sortie (ptks/s et/ou bytes/s)
  * TZSP over IPv6
  * getopts -v et --help
  * Licence GPL
*/

/* Functions declaration */
void capture_loop(char pcap_filter[]); 
void process_packet(u_char *void_args, const struct pcap_pkthdr* pkthdr, const u_char * packet);

/* Custom types definition */
typedef struct _process_packet_args {
	/* Types like in libpcap because I don't want to mess around hours */
	u_int captured_pkt_count;
	u_int sent_pkt_count;
	int udp_socket;
	struct sockaddr_in udp_sockaddr;
} process_packet_args_t;

/* Constants */
const int sockaddr_len=sizeof(struct sockaddr_in);


/* Flags from commandline parsing */
static int opt_verbose;
/* Option arguments from commandline parsing */
static char *opt_iface=NULL;
static char *opt_host=NULL;
static char *opt_port=NULL;
static char *opt_snaplen=NULL;

pcap_t *pcap_handle = NULL; 

void usage(char progname[]) {
	printf("Usage : %s [--verbose] [--brief] [-i <iface>] [-h <host>] [-p <port>] [custom_pcap_filter]\n", progname);
	printf("\t<iface> : Interface name to capture from (Default : first available interface)\n");
	printf("\t<host> : Host (or IPv4 address) for sending captured packet headers (Default : '%s')\n", DEFAULT_HOST);
	printf("\t<port> : Port for sending captured packet headers (Default '%s')\n", DEFAULT_PORT);
	printf("\t<custom_pcap_filter> : libpcap capture filter (Default '%s')\n", DEFAULT_FILTER);
	exit(1);
}

void sig_handler(int signo) {
	static int pcap_break=0;

	switch (signo) {
		case SIGINT:
			if (pcap_break==0) {
				pcap_break=1;
				//fprintf(stderr, "DEBUG : pcap_breakloop(pcap_handle);\n");
				pcap_breakloop(pcap_handle);
			}
			break;
		default:
			fprintf(stderr, "Catched an unhandled signal : %i\n", signo);
	}
}

int main(int argc, char *argv[]) {

	/* Command line args parsing */
	int c,i,j,pcap_filter_len;
	char *pcap_filter=NULL;

	while (1) {
		static struct option long_options[] = {
			/* These options set a flag. */
			{"verbose", no_argument, &opt_verbose, 1},
			{"brief",   no_argument, &opt_verbose, 0},
			/* These options don't set a flag.
			   We distinguish them by their indices. */
			{"interface", required_argument, 0, 'i'},
			{"host",      required_argument, 0, 'h'},
			{"port",      required_argument, 0, 'p'},
			{"snaplen",   required_argument, 0, 's'},
			{0, 0, 0, 0}
		};
		/* getopt_long stores the option index here. */
		int option_index = 0;

		c = getopt_long(argc, argv, "i:h:p:s:", long_options, &option_index);

		/* Detect the end of the options. */
		if (c == -1) break;

		switch (c) {
			case 0:
				/* If this option set a flag, do nothing else now. */
				if (long_options[option_index].flag != 0) break;

				//printf ("option %s", long_options[option_index].name);
				//if (optarg) printf (" with arg %s", optarg);
				//printf ("\n");

				break;
			case 'i': opt_iface=	strdup(optarg); break;
			case 'h': opt_host=	strdup(optarg); break;
			case 'p': opt_port=	strdup(optarg); break;
			case 's': opt_snaplen=	strdup(optarg); break;
			case '?':
				/* getopt_long already printed an error message. */
				usage(basename(argv[0]));
				break;
			default:
				abort();
		}
	}

	// Assign default value if no user supplied value
	if (opt_host==NULL)	opt_host=strdup(DEFAULT_HOST);
	if (opt_port==NULL)	opt_port=strdup(DEFAULT_PORT);
	if (opt_snaplen==NULL)	opt_snaplen=strdup(DEFAULT_SNAPLEN);
	// No default for opt_iface because will be choosen at runtime

	//if (verbose_flag) puts ("verbose flag is set");

	/* Construct the pcap_filter */
	pcap_filter_len=0;
	if (optind < argc) {
		// Any remaining command line arguments is for pcap_filter
		for (i=optind; i<argc; i++) pcap_filter_len += strlen(argv[i])+1;

		pcap_filter=calloc(pcap_filter_len+1, sizeof(char));
		if (pcap_filter==NULL) {
			fprintf(stderr,"calloc() failed for pcap_filter\n");
			abort();
		}

		for (i=optind, j=0; i<argc; i++) {
			if (i<argc-1) {
				 j+=sprintf(pcap_filter+j, "%s ", argv[i]);
			} else {
				j+=sprintf(pcap_filter+j, "%s", argv[i]);
			}
		}
	} else {
		// Else, use the default filter (skip the TZSP packet flow to host:port)
		pcap_filter_len = strlen(DEFAULT_FILTER) -4 + strlen(opt_host) + strlen(opt_port) + 1;
		pcap_filter=calloc(pcap_filter_len+1, sizeof(char));
		j=sprintf(pcap_filter, DEFAULT_FILTER, opt_host, opt_port);
	}
	//printf ("Capture filter will be : '%s' (pcap_filter_len-j==%i)\n", pcap_filter, pcap_filter_len-j);


	signal(SIGINT, sig_handler);
	/* Run the capture loop */
	capture_loop(pcap_filter);

	free(pcap_filter);

	free(opt_host);
	free(opt_port);
	free(opt_snaplen);
	if (opt_iface!=NULL) free(opt_iface);

	return 0;
}


void capture_loop(char pcap_filter[]) { 

	//Global for signal handling
	//pcap_t *pcap_handle = NULL; 
	static char *pcap_device=NULL;
	char pcap_errbuf[PCAP_ERRBUF_SIZE];
	struct bpf_program pcap_filter_prog;
	struct pcap_stat stat;
	process_packet_args_t process_packet_args;
	int snaplen;

	memset(pcap_errbuf,0,PCAP_ERRBUF_SIZE); 
	memset(&process_packet_args,0,sizeof(process_packet_args_t)); 

	if (opt_iface == NULL) {
		/* Get the name of the first device suitable for capture */ 
		if ( (pcap_device = pcap_lookupdev(pcap_errbuf)) == NULL){
			fprintf(stderr, "ERROR: %s\n", pcap_errbuf);
			exit(20);
		}
	} else {
		/* Use user-suplied interface */
		pcap_device=opt_iface;
	}

	snaplen=atoi(opt_snaplen);
	if ( snaplen < MIN_BYTES_TO_CAPTURE || snaplen > MAX_BYTES_TO_CAPTURE ) {
		fprintf(stderr, "snaplen of %i is not in the allowed range of [%i..%i]\n", snaplen, MIN_BYTES_TO_CAPTURE, MAX_BYTES_TO_CAPTURE);
		exit(12);
	}

	/* Already done : memset((char *) &(process_packet_args.udp_sockaddr), 0, sockaddr_len); */
	process_packet_args.udp_sockaddr.sin_family = AF_INET;
	process_packet_args.udp_sockaddr.sin_port = htons(atoi(opt_port));
	if (inet_aton(opt_host, &(process_packet_args.udp_sockaddr.sin_addr))==0) {
		fprintf(stderr, "inet_aton : Impossible de parser l'IP %s\n",opt_host);
		exit(11);
	}
	

	/* Initialisation pcap */
	if (opt_verbose) printf("Opening device %s for capturing\n", pcap_device);

	/* Open device in promiscuous mode */ 
	if ( (pcap_handle = pcap_open_live(pcap_device, snaplen, 1,  512, pcap_errbuf)) == NULL){
		fprintf(stderr, "ERROR: %s\n", pcap_errbuf);
		exit(21);
	}

	if (opt_verbose) printf("Compiling the following capture filter : '%s'\n", pcap_filter);
	if ( pcap_compile(pcap_handle,&pcap_filter_prog,pcap_filter,1,PCAP_NETMASK_UNKNOWN) == -1 ) {
		fprintf(stderr,"ERROR : %s\n", pcap_geterr(pcap_handle) );
		exit(22);
	}

	if ( pcap_setfilter(pcap_handle,&pcap_filter_prog) == -1 ) {
		fprintf(stderr,"ERROR : %s\n", pcap_geterr(pcap_handle) );
		exit(23);
	}

	pcap_freecode(&pcap_filter_prog);

	/* Initialisation socket UDP */
	process_packet_args.udp_socket=socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
	if (process_packet_args.udp_socket==-1) {
		perror("socket");
		exit(10);
	}

	/* Loop forever & call process_packet() for every received packet */ 
	if ( pcap_loop(pcap_handle, -1, process_packet, (u_char *)&process_packet_args) == -1) {
		fprintf(stderr, "ERROR: %s\n", pcap_geterr(pcap_handle) );
		exit(25);
	}

	fprintf(stderr, "\n%u packets captured\n%u TZSP packets sent\n",
		 process_packet_args.captured_pkt_count, process_packet_args.sent_pkt_count);
	
	if (pcap_stats(pcap_handle, &stat) == -1) {
		fprintf(stderr, "ERROR: %s\n", pcap_geterr(pcap_handle) );
		exit(26);
	}
	fprintf(stderr, "%u packets received by filter\n%u packets dropped by kernel\n",
		stat.ps_recv, stat.ps_drop);

	pcap_close(pcap_handle);
	close(process_packet_args.udp_socket);
}


/* process_packet(): Callback function called by pcap_loop() everytime a packet */
/* arrives to the network card. This function prints the captured raw data in  */
/* hexadecimal.                                                                */
void process_packet(u_char *void_args, const struct pcap_pkthdr* pkthdr, const u_char * packet) {

	static time_t old_tv_sec=0;
	static u_int old_captured_pkt_count=0;

	int res;
	double throughput;

	uint32_t ts; /* In network byte order */
	//uint32_t pkt_num; /* In network byte order */
	uint16_t len; /* In network byte order */
	char buf[UDP_SEND_BUFLEN];
	/* struct timespec ts_reqsleep; For simulation */

	process_packet_args_t *args=(process_packet_args_t *)void_args; 

	/* Catpured packet counting and displaying (max once by second) */
	args->captured_pkt_count++;

	if (old_tv_sec != pkthdr->ts.tv_sec) {
		/*printf("DEBUG : throughput=((double)  %u - %u ) / (  %u - %u )\n",
			args->captured_pkt_count, old_captured_pkt_count, pkthdr->ts.tv_sec, old_tv_sec);
		*/
		throughput=((double) args->captured_pkt_count - old_captured_pkt_count ) / (pkthdr->ts.tv_sec - old_tv_sec);
		printf("\rPacket Count: %u (%8.1f pkt/s)", args->captured_pkt_count, throughput);
		fflush(stdout);
		old_tv_sec=pkthdr->ts.tv_sec;
		old_captured_pkt_count= args->captured_pkt_count;
	}

	//TODO : bwlimit HERE (not before, no after)

	/* Variable fields for TZSP packet */
	ts=htonl((uint32_t) pkthdr->ts.tv_sec); /* TODO : this cast is bullshit ? */
	len=htons((uint16_t) pkthdr->len);
	//pkt_num=htons((uint32_t) args->captured_pkt_count);

	/* TaZmen Sniffing Protocol (TZSP) packet forging */

	//memset((char *) &buf, 0, UDP_SEND_BUFLEN); /* Buffer reset not useful for now */
	buf[0]=0x01; /* Version */
	buf[1]=0x01; /* Type == Packet for transmit */
	buf[2]=0x00; /* Encapsuled protocol (2 bytes) 0x0001 == Ethernet */
	buf[3]=0x01;

	buf[4]=0x0D; /* Tag type == TAG_TIMESTAMP */
	buf[5]=0x04; /* Tag length == 4 bytes */

	/* buf[6,7,8,9] Timestamp on 4 bytes (network order) */
	memcpy(buf+6, &ts, 4);	

	/* Wireshark don't dissect that */
	//buf[10]=0x28; /* Tag type TAG_PACKET_COUNT */
	//buf[11]=0x04; /* Tag length == 4 bytes */
	//memcpy(buf+12,&pkt_num, 4);

	buf[10]=0x29; /* Tag type TAG_RX_FRAME_LENGTH */
	buf[11]=0x02; /* Tag length == 2 bytes */
	memcpy(buf+12,&len, 2);

	buf[14]=0x00; /* Tag type TAG PADDING (for alignement) */
	buf[15]=0x01; /* Tag type TAG END */

	/* Raw packet copy */
	//TODO : assert that pkthdr->caplen < MAX_TZSP_PAYLOAD
	memcpy(buf+16,packet, pkthdr->caplen);

	/* TZSP packet sending (over the UDP socket) */
	res=sendto(args->udp_socket, buf, 16+pkthdr->caplen, 0, (struct sockaddr *)&(args->udp_sockaddr), sockaddr_len);
	if (res==-1) {
		fprintf(stderr, "sendto() error\n");
	} else {
		args->sent_pkt_count++;
	}

	/* Slow sending simulation
	ts_reqsleep.tv_sec=0;
	ts_reqsleep.tv_nsec=10000;
	nanosleep(&ts_reqsleep, NULL);
	*/
}