/*

    Pcap2tzsp allows to capture ethernet trafic and send all headers
    to a defined host using the TaZmen Sniffing Protocol (TZSP).

    Copyright (C) 2012 Ludovic Pouzenc <lpouzenc@gmail.com>

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.


    This file is part of Pcap2tzsp. See LICENCE file.
*/

#include "pcap2tzsp.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 "localhost"
#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)

/* Max len of a numerical form of host address 39 chars in IPv6 */
#define NI_MAXHOST_NUMERIC 40

/*
TODO List
  * Essayer  WIN32_LEAN_AND_MEAN avec VS
  * Implémenter une bwlimit en sortie (ptks/s et/ou bytes/s)
  * Calcul dynamique de MAX_TZSP_PAYLOAD
*/

/* Functions declaration */
SOCKET make_dgram_socket(char host[], char service[], int hint_flags, char **resolved_address);
void start_capture_loop(char *pcap_filter);
void process_packet(u_char *void_args, const struct pcap_pkthdr* pkthdr, const u_char * packet);
void make_tzsp_notags(u_char *out_tzsp, size_t *tzsplen, const u_char *pdu, size_t pdu_len);
void make_tzsp_basic(u_char *out_tzsp, size_t *tzsplen, const u_char *pdu, size_t pdu_len, time_t original_packet_timestamp, size_t original_packet_len );


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


/* 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>] [-s <snaplen>] [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<snaplen> : Snarf  snaplen  bytes  of  data from each packet (Default '%s')\n", DEFAULT_SNAPLEN); 
	printf("\t<custom_pcap_filter> : libpcap capture filter (Default '%s')\n", DEFAULT_FILTER);
	exit(1);
}

#ifndef COMPAT_LIBPCAP_0_7
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);
	}
}
#endif

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

#ifdef WIN32
	// Winsock2 mystic stuff
	WSADATA wsaData = {0};
	int iResult;
#endif

	/* Command line args parsing */
	int c,i,j;
	size_t pcap_filter_len; /* size_t is "mandatory" for WIN32 */
	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. */
			{"help",    no_argument, 0, 'H'},
			{"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:v", 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 'v': opt_verbose=	1; break;
			case 'H':
			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+=my_snprintf(pcap_filter+j,pcap_filter_len-j, "%s ", argv[i]);
			} else {
				j+=my_snprintf(pcap_filter+j,pcap_filter_len-j, "%s", argv[i]);
			}
		}
	}


#ifdef WIN32
	// Initialize Winsock
	iResult = WSAStartup(MAKEWORD(2, 2), &wsaData);
    if (iResult != 0) {
        fprintf(stderr, "WSAStartup failed: %d\n", iResult);
        return -10;
    }
#endif

#ifndef COMPAT_LIBPCAP_0_7
	signal(SIGINT, sig_handler);
#endif
	/* Run the capture loop */
	start_capture_loop(pcap_filter);

#ifdef WIN32
	WSACleanup();
#endif
	free(opt_host);
	free(opt_port);
	free(opt_snaplen);
	if (opt_iface!=NULL) free(opt_iface);
	if (pcap_filter!=NULL) free(pcap_filter);

	return 0;
}

SOCKET make_dgram_socket(char host[], char service[], int hint_flags, char **resolved_address) {

	/* Code taken from Client program example of getaddrinfo(3) manpage
	   Tweaks added for WIN32 "compatibility"
	*/

	struct addrinfo hints;
	struct addrinfo *result, *rp;
	char hbuf[NI_MAXHOST_NUMERIC];
	int s;
	SOCKET sfd = INVALID_SOCKET;

	/* Obtain address(es) matching host/port */

	memset(&hints, 0, sizeof(struct addrinfo));
	hints.ai_family = AF_UNSPEC;	/* Allow IPv4 or IPv6 */
	hints.ai_socktype = SOCK_DGRAM;	/* Datagram socket */
	hints.ai_flags = hint_flags | AI_CANONNAME;
	hints.ai_protocol = 0;		/* Any protocol */

	s = getaddrinfo(host, service, &hints, &result);
	if (s != 0) {
		fprintf(stderr, "getaddrinfo: %s\n", gai_strerror(s));
		return -2;
	}

	/* getaddrinfo() returns a list of address structures.
	   Try each address until we successfully connect(2).
	   If socket(2) (or connect(2)) fails, we (close the socket
	   and) try the next address. */

	for (rp = result; rp != NULL; rp = rp->ai_next) {
		sfd = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
		if (sfd == INVALID_SOCKET)
			continue;

		if (connect(sfd, rp->ai_addr, rp->ai_addrlen) != -1)
			break;		/* Success */
#ifdef WIN32
		closesocket(sfd);
#else
		close(sfd);
#endif
	}

	if (rp == NULL) {		/* No address succeeded */
		freeaddrinfo(result);
		return -1;
	}

	/* This part was not in the example as is */
	if ( resolved_address != NULL) {
		/* If wanted, return numerical form of the choosen host address */
		if ( getnameinfo(rp->ai_addr, rp->ai_addrlen, hbuf,
			NI_MAXHOST_NUMERIC, NULL, 0, NI_NUMERICHOST) == 0 ) {
			*resolved_address=strdup(hbuf);
		}
	}

	freeaddrinfo(result);		/* No longer needed */
	return sfd;
}


void start_capture_loop(char *pcap_filter) {

	//Global for signal handling
	//pcap_t *pcap_handle = NULL;
	char *resolved_address=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, hints_flags, pcap_filter_synthetised=0;
	size_t pcap_filter_len; /* size_t is "mandatory" for WIN32 */

	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 (maybe you have to start this program as root)\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);
	}

	/* UDP socket creation */
	hints_flags=AI_NUMERICSERV | AI_ADDRCONFIG | AI_V4MAPPED;
	if (opt_verbose) printf("Opening socket for sending TZSP to %s:%s\n", opt_host, opt_port);
	if ( ( process_packet_args.socket = make_dgram_socket(opt_host, opt_port, hints_flags, &resolved_address) ) < 0 ) {
		fprintf(stderr, "ERROR : Couldn't make a socket to %s:%s\n", opt_host, opt_port);
		exit(10);
	}
	if (opt_verbose) printf("Socket opened with the following destination : %s:%s\n", resolved_address, opt_port);

	/* pcap init */
	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 (pcap_filter==NULL) {
		pcap_filter_synthetised=1;
		// Synthethise a default filter that skip the TZSP packet flow to host:port
		pcap_filter_len = strlen(DEFAULT_FILTER) -4 + strlen(resolved_address) + strlen(opt_port) + 1;
		pcap_filter=calloc(pcap_filter_len, sizeof(char));
		//XXX pos=snprintf(pcap_filter, pcap_filter_len, DEFAULT_FILTER, resolved_address, opt_port);
		/*res=*/ (void)my_snprintf(pcap_filter, pcap_filter_len, DEFAULT_FILTER, resolved_address, opt_port);
		//printf ("DEBUG : Capture filter will be : '%s' (pcap_filter_len-res==%i)\n", pcap_filter, pcap_filter_len-res);
	}
	free(resolved_address);

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

	if (pcap_filter_synthetised) free(pcap_filter);

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

	/* 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);
#ifdef WIN32
	closesocket(process_packet_args.socket);
#else
	close(process_packet_args.socket);
#endif
}


/* process_packet(): Callback function called by pcap_loop()
   everytime a packet arrives to the network card. */
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;
	size_t tzsp_len;
	double throughput;

	u_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, not after)

	/* TaZmen Sniffing Protocol (TZSP) packet forging */
	tzsp_len=UDP_SEND_BUFLEN;
	switch(0) { /* TODO */
		case 1:
			make_tzsp_notags(buf, &tzsp_len,packet,pkthdr->caplen);
			break;
		default:			
			make_tzsp_basic(buf, &tzsp_len,packet,pkthdr->caplen, pkthdr->ts.tv_sec, pkthdr->len);
	}
	if (tzsp_len < 0) {
		fprintf(stderr, "make_tzsp_*() problem\n");
		return;
	}

	/* TZSP packet sending (over the UDP socket) */

	/* Slow sending simulation
	ts_reqsleep.tv_sec=0;
	ts_reqsleep.tv_nsec=10000;
	nanosleep(&ts_reqsleep, NULL);
	*/
#ifdef WIN32
	res=send(args->socket, buf, tzsp_len, 0);
#else
	res=write(args->socket, buf, tzsp_len);
#endif
	if (res != tzsp_len) {
		fprintf(stderr, "write() on UDP socket error (written %i of %zu bytes)\n", res, tzsp_len);
		return;
	}

	args->sent_pkt_count++;
}

void make_tzsp_notags(u_char *out_tzsp, size_t *tzsplen, const u_char *pdu, size_t pdu_len) {
	const char tzsp_header_tpl[] = {
		0x01,		/* Version 1 */
		0x01,		/* Type : Packet for transmit */
		0x00, 0x01,	/* Encapsulated protocol : 0x001 == Ethernet */
		0x01		/* Tag type : TAG END */
	};
	const size_t tzsp_header_len = sizeof(tzsp_header_tpl);
	size_t needed = tzsp_header_len + pdu_len;

	/* Check if enought space in out buffer */
	if ( *tzsplen < needed ) {
		*tzsplen=-1;
		return;
	}

	/* Header copy */
	memcpy(out_tzsp, tzsp_header_tpl, tzsp_header_len);

        /* Raw packet copy */
        memcpy(out_tzsp + tzsp_header_len, pdu, pdu_len);

	/* Update tzsplen to the effecive len */
	*tzsplen=needed;

	return;
}

void make_tzsp_basic(u_char *out_tzsp, size_t *tzsplen, const u_char *pdu, size_t pdu_len, time_t original_packet_timestamp, size_t original_packet_len ) {
	const char tzsp_header_tpl[] = {
		0x01,		/* Version 1 */
		0x01,		/* Type : 0x01 == Packet for transmit */
		0x00, 0x01,	/* Encapsulated protocol : 0x0001 == Ethernet */

		0x0D,		/* Tag type TAG_TIMESTAMP */
		0x04,		/* Tag len : 4 bytes */
		0x00, 0x00, 0x00, 0x00, /* Tag Value (placeholder) */

		0x29,		/* Tag type TAG_RX_FRAME_LENGTH */
		0x02,		/* Tag len : 2 bytes */
		0x00, 0x00, 	/* Tag Value (placeholder) */

		0x00,		/* Tag type : TAG PADDING (for alignement) */
		0x01		/* Tag type : TAG END */
	};
	const size_t tzsp_header_len = sizeof(tzsp_header_tpl);
	uint32_t tag_timestamp_value;
	uint16_t tag_packet_len_value;
	size_t needed;

	/* Check if enough space in out buffer */
	needed = tzsp_header_len + pdu_len;
	if ( *tzsplen < needed ) {
		*tzsplen=-1;
		return;
	}

	/* "Compute" tags values (network byte order) */
	tag_timestamp_value=htonl((uint32_t) original_packet_timestamp);
	tag_packet_len_value=htons((uint16_t) original_packet_len);

	/* Header template copy */
	memcpy(out_tzsp, tzsp_header_tpl, tzsp_header_len);

	/* Tag values copy */
	memcpy(out_tzsp+6, &tag_timestamp_value, sizeof(uint32_t));
	memcpy(out_tzsp+12, &tag_packet_len_value, sizeof(uint16_t));

        /* Raw packet copy */
        memcpy(out_tzsp + tzsp_header_len, pdu, pdu_len);

	/* Update tzsplen to the effecive len */
	*tzsplen=needed;

	return;
}