/* A much more complicated libtrace program designed to demonstrate combining
 * various elements of libtrace to create a useful tool.
 *
 * Specifically, this program calculates the amount of header overhead for
 * TCP and UDP traffic compared with the amount of application payload. It
 * writes the byte counts regularly to generate data suitable for a time series
 * graph.
 *
 */
#include "libtrace.h"
#include <stdio.h>
#include <inttypes.h>
#include <err.h>
#include <assert.h>
#include <getopt.h>
#include <string.h>
#include <stdlib.h>

uint64_t udp_header = 0;
uint64_t udp_payload = 0;
uint64_t tcp_header = 0;
uint64_t tcp_payload = 0;
uint64_t not_ip = 0;

uint32_t next_report = 0;
uint32_t interval = 10;		/* Reporting interval defaults to 10 seconds. */

/* This enum defines values for all the possible protocol cases that this
 * program is interested in */
typedef enum {
	DEMO_PROTO_TCP,		/* The packet is a TCP packet */
	DEMO_PROTO_UDP,		/* The packet is a UDP packet */
	DEMO_PROTO_NOTIP,	/* The packet is NOT an IP packet */
	DEMO_PROTO_OTHER,	/* The packet is none of the above */
	DEMO_PROTO_UNKNOWN	/* Haven't yet determined anything about the
				   packet */
} demo_proto_t;

static void print_stats() {
	printf("%u,%" PRIu64 ",%" PRIu64 ",%" PRIu64 ",%" PRIu64 ",%" PRIu64 "\n",
		next_report, tcp_header, tcp_payload, udp_header, 
		udp_payload, not_ip);
}

static void check_report(libtrace_packet_t *packet) {
	struct timeval ts;

	/* Get the timestamp for the current packet */
	ts = trace_get_timeval(packet);

	/* If next_report is zero, then this is the first packet from the
	 * trace so we need to determine the time at which the first report 
	 * must occur, i.e. "interval" seconds from now. */

	if (next_report == 0) {
		next_report = ts.tv_sec + interval;

		/* Good opportunity to print some column headings */
		printf("Time,TCP Headers,TCP Payload,UDP Headers,UDP Payload,Not IP\n");
	}

	/* Check whether we need to report our stats
	 *
	 * Compare the timestamp for the current packet against the time that
	 * the next report is due, a la timedemo.c
	 */

	while ((uint32_t)ts.tv_sec > next_report) {
		/* Print all our stats */
		print_stats();

		/* Reset the counters */
		tcp_header = 0;
		tcp_payload = 0;
		udp_header = 0;
		udp_payload = 0;
		not_ip = 0;

		/* Determine when the next report is due */
		next_report += interval;
	}
}

/* Calculates the number of bytes consumed by meta-data headers such as
 * Linux SLL, RadioTap, etc.
 */
static uint64_t calc_meta_size(libtrace_packet_t *packet) {

	uint64_t meta_size = 0;
	void *meta_ptr = NULL;
	libtrace_linktype_t ltype;
	uint32_t remaining;
	uint32_t prev_rem;

	/* Get a pointer to the meta-data header */
	meta_ptr = trace_get_packet_meta(packet, &ltype, &remaining);

	/* If the result is NULL, there are no meta-data headers present */
	if (meta_ptr == NULL) 
		return meta_size;


	/* Skip over any subsequent meta-data headers */
	while (remaining > 0) {
		prev_rem = remaining;
		void *nexthdr = trace_get_payload_from_meta(meta_ptr,
			&ltype, &remaining);

		/* If nexthdr is NULL, the current header is NOT a meta-data
		 * header (and is almost certainly a link layer header) */
		if (nexthdr == NULL) 
			break;
		
		/* Sanity check as remaining should never get larger! */
		assert(prev_rem >= remaining);

		/* Otherwise the header we called get_payload on was a 
		 * meta-data header so we need to add its length to the total
		 * meta header size */	
		meta_size += (prev_rem - remaining);

		/* Prepare ourselves for the next pass through the loop */
		meta_ptr = nexthdr;
	}

	return meta_size;


}

/* Calculates the number of bytes consumed by link layer headers. Note that
 * this will include any "layer 2.5" headers such as MPLS, VLAN or PPP.
 */
static uint64_t calc_link_size(libtrace_packet_t *packet) {
	void *link_ptr;
	void *nexthdr;
	libtrace_linktype_t linktype;
	uint32_t remaining;
	uint32_t prev_rem;
	uint16_t ethertype;

	uint64_t link_size = 0;

	/* Start by finding the layer 2 header */
	link_ptr = trace_get_layer2(packet, &linktype, &remaining);

	/* If there is no layer 2 header, the total link layer has to be
	 * zero bytes in size */
	if (link_ptr == NULL) 
		return link_size;

	/* Calculate the size of the first layer 2 header by comparing
	 * remaining before and after we call trace_get_payload_from_layer2
	 */
	prev_rem = remaining;
	nexthdr = trace_get_payload_from_layer2(link_ptr, linktype, 
			&ethertype, &remaining);

	/* Sanity check - remaining should never get larger! */
	assert(prev_rem >= remaining);
	/* Add the size of the layer 2 header to our overall link layer size */
	link_size += (prev_rem - remaining);

	/* Skip over any layer 2.5 headers, adding their size to our total
	 * link layer size.  */
	while (remaining > 0) {
		if (nexthdr == NULL)
			break;
		prev_rem = remaining;

		/* Ethertype will always contain the type of the current
		 * header that we are up to, thanks to the efforts of the
		 * trace_get_payload_from_* functions */
		switch(ethertype) {
			case 0x8100: 	/* VLAN */
				nexthdr = trace_get_payload_from_vlan(
					nexthdr, &ethertype, &remaining);
				break;
			case 0x8847:	/* MPLS */
				nexthdr = trace_get_payload_from_mpls(
					nexthdr, &ethertype, &remaining);
				break;
			case 0x8864:	/* PPPoE */
				/* This will also skip the PPP header */
				nexthdr = trace_get_payload_from_pppoe(
					nexthdr, &ethertype, &remaining);
				break;
			default:
				/* This is just to provide a stopping condition
				 * for the while loop. */
				nexthdr = NULL;			
		}
		
		/* If we have reached a non-layer 2.5 header, i.e. IP, we 
		 * want to fall out and return the total size */
		if (nexthdr == NULL)
			break;
		
		/* Otherwise, add the length of the skipped header to the
		 * total, being sure to perform our usual sanity check first */
		assert(prev_rem >= remaining);
		link_size += (prev_rem - remaining);
		
	}

	/* Return the total link layer size */
	return link_size;
}

/* Calculates the number of bytes consumed by IP headers, including IPv6 */
static uint64_t calc_ip_size(libtrace_packet_t *packet, demo_proto_t *proto) {

	uint64_t ip_size = 0;
	
	void *ip_hdr;
	void *nexthdr = NULL;;
	uint16_t ethertype;
	uint8_t protocol;
	uint32_t remaining;
	uint32_t prev_rem;
	libtrace_ip6_t *ip6;

	/* Start by finding the first layer 3 header */
	ip_hdr = trace_get_layer3(packet, &ethertype, &remaining);

	/* If no layer 3 headers are present, be sure to set proto 
	 * appropriately so the total header length is added to the right
	 * category */
	if (ip_hdr == NULL) {
		*proto = DEMO_PROTO_NOTIP;
		return ip_size;
	}

	prev_rem = remaining;

	/* Unlike at the link layer, there is less scope for endlessly stacked
	 * headers so we don't need a fancy while loop */

	/* Remember, ethertype tells us the type of the layer 3 header so we
	 * can cast appropriately */
	switch(ethertype) {
		case 0x0800:	/* IPv4 */
			/* Skip past the IPv4 header */
			nexthdr = trace_get_payload_from_ip(ip_hdr, &protocol, 
					&remaining);
			
			/* Check for v6 over v4 and skip over it if present */
			if (nexthdr && protocol == 41) {
				ip6 = (libtrace_ip6_t *)nexthdr;
				nexthdr = trace_get_payload_from_ip6(
					ip6, &protocol, &remaining);
			}
			break;
		case 0x86DD:	/* IPv6 */
			/* Skip past the IPv6 header */
			ip6 = (libtrace_ip6_t *)ip_hdr;
			nexthdr = trace_get_payload_from_ip6(ip6, &protocol, 
					&remaining);
			break;
		default:
			/* Somehow we managed to get a layer 3 header that is
			 * neither v4 nor v6 */
			*proto = DEMO_PROTO_NOTIP;
			return ip_size;
	}

	/* Update our layer 3 size with the number of bytes we just skipped
	 * past */
	assert(prev_rem >= remaining);
	ip_size += (prev_rem - remaining);

	/* We can also use the protocol value from the get_payload function
	 * to determine the transport layer protocol */
	if (protocol == 6) 
		*proto = DEMO_PROTO_TCP;
	else if (protocol == 17)
		*proto = DEMO_PROTO_UDP;
	else
		*proto = DEMO_PROTO_OTHER;
	
	/* Return our total layer 3 size */
	return ip_size;

}

/* Calculates the number of bytes consumed by the transport header, including
 * options etc. */
static uint64_t calc_transport_size(libtrace_packet_t *packet) {
	
	uint64_t trans_size = 0;

	void *transport;
	void *nexthdr;
	uint8_t proto;
	uint32_t remaining;
	uint32_t prev_rem;
	libtrace_tcp_t *tcp;
	libtrace_udp_t *udp;

	/* Start by finding the transport header */
	transport = trace_get_transport(packet, &proto, &remaining);

	/* No transport header makes our life very easy - we can just return
	 * zero */
	if (transport == NULL)
		return trans_size;

	prev_rem = remaining;

	/* Skip past the transport header. Transport headers (at least the ones
	 * we're interested in) can't be stacked so we only ever need to skip
	 * past the one header */

	/* Switch based on the protocol value set by trace_get_transport */
	switch (proto) {
		case 6:		/* TCP */
			tcp = (libtrace_tcp_t *)transport;
			nexthdr = trace_get_payload_from_tcp(tcp, &remaining);
			break;
		case 17:	/* UDP */
			udp = (libtrace_udp_t *)transport;
			nexthdr = trace_get_payload_from_udp(udp, &remaining);
			break;
		default:
			/* We have no interest in ICMP, GRE etc, and we 
			 * should never have entered this function if the
			 * packet is using those protocols anyway! */
			fprintf(stderr, "Unexpected protocol: %u\n", proto);
			return 0;
	}

	/* If we don't have any post-transport payload, just return the 
	 * transport header size */
	if (!nexthdr)
		return trans_size;
	
	/* Determine how many bytes we just skipped over and add it to the
	 * total transport size */
	assert(prev_rem >= remaining);
	trans_size += (prev_rem - remaining);

	/* Return the total size */
	return trans_size;
}

static uint64_t calc_header_size(libtrace_packet_t *packet, demo_proto_t *proto) {

	uint64_t size = 0;

	/* Start with any meta-data headers */
	size += calc_meta_size(packet);

	/* Work out the size of link layer headers */
	size += calc_link_size(packet);

	/* Determine the size of the IP headers */
	size += calc_ip_size(packet, proto);

	/* If the previous function call determined we were not an IP packet,
	 * we can drop out now and return the current size */
	if (*proto == DEMO_PROTO_NOTIP) 
		return size;

	/* We can also drop out if the packet is not using a protocol that we
	 * are interested in */
	if (*proto == DEMO_PROTO_OTHER || *proto == DEMO_PROTO_UNKNOWN)
		return 0;

	/* Add on the transport headers */
	size += calc_transport_size(packet);

	/* Return the total amount of headers */
	return size;
}

static uint64_t calc_payload_size(libtrace_packet_t *packet, demo_proto_t proto)
{

	uint64_t ip_plen = 0;
	uint64_t headers = 0;

	void *layer3;
	uint16_t ethertype;
	uint32_t remaining;

	layer3 = trace_get_layer3(packet, &ethertype, &remaining);

	/* This should NEVER happen, but it's a good habit to check it anyway */
	if (layer3 == NULL)
		return 0;

	/* Find the payload length in the IP header
	 *
	 * We also determine the size of the IP header (again!) as the payload
	 * length includes the IP and transport headers */
	
	if (ethertype == 0x0800) {   /* IPv4 header */
		libtrace_ip_t *ip = (libtrace_ip_t *)layer3;
		
		/* Remember to byte swap! */
		ip_plen = ntohs(ip->ip_len);
		/* This value is only 4 bits so byteswapping is unnecessary */
		headers += (4 * ip->ip_hl);
	
	
	} else if (ethertype == 0x86DD) { /* IPv6 header */
		libtrace_ip6_t *ip = (libtrace_ip6_t *)layer3;
		
		/* Remember to byte swap! */
		ip_plen = ntohs(ip->plen);
		/* IPv6 does not have a variable length header */
		headers += sizeof(libtrace_ip6_t);
	
	
	} else {
		/* Not an IP packet - this should also never happen */
		return 0;
	}

	
	/* Now we need to subtract the size of the transport header from the
	 * IP payload length. */
	if (proto == DEMO_PROTO_TCP) {
		
		/* Determine the size of the TCP header so we can subtract
		 * that from our total payload length */
		
		/* Since I already know the protocol and only need to 
		 * access a single value inside the TCP header, I can use
		 * the trace_get_tcp() helper function instead of the more
		 * verbose trace_get_transport() . */
		libtrace_tcp_t *tcp = trace_get_tcp(packet);
		if (tcp == NULL)
			return 0;

		/* Again, byteswapping not required because the doff field is
		 * only a single byte in size*/
		headers += (tcp->doff * 4);
	}

	if (proto == DEMO_PROTO_UDP) {
		/* UDP has a fixed length header so we don't even need to use
		 * trace_get_udp() */
		headers += sizeof(libtrace_udp_t);
	}

	assert(headers <= ip_plen);

	/* Subtract the length of the IP and transport headers from the 
	 * payload length contained within the IP header */
	return ip_plen - headers;

}

static void per_packet(libtrace_packet_t *packet)
{
	uint64_t header_size = 0;
	uint64_t payload_size = 0;
	demo_proto_t protocol = DEMO_PROTO_UNKNOWN;

	/* Check if we're due to report some stats */
	check_report(packet);

	/* We need to determine the amount of header in this packet */
	header_size = calc_header_size(packet, &protocol);
	
	/* Now determine the payload size, if necessary */
	if (protocol == DEMO_PROTO_TCP || protocol == DEMO_PROTO_UDP) {
		payload_size = calc_payload_size(packet, protocol);
	}

	/* Update the appropriate counters */
	switch(protocol) {
		case DEMO_PROTO_TCP:
			tcp_header += header_size;
			tcp_payload += payload_size;
			break;

		case DEMO_PROTO_UDP:
			udp_header += header_size;
			udp_payload += payload_size;
			break;

		case DEMO_PROTO_NOTIP:
			not_ip += header_size;
			break;

		case DEMO_PROTO_OTHER:
			break;

		case DEMO_PROTO_UNKNOWN:
			break;
	}


}

/* Due to the amount of error checking required in our main function, it
 * is a lot simpler and tidier to place all the calls to various libtrace
 * destroy functions into a separate function.
 */
static void libtrace_cleanup(libtrace_t *trace, libtrace_packet_t *packet, 
		libtrace_filter_t *filter) {
	
	/* It's very important to ensure that we aren't trying to destroy
	 * a NULL structure, so each of the destroy calls will only occur
	 * if the structure exists */
	if (trace)
		trace_destroy(trace);
	
	if (packet)
		trace_destroy_packet(packet);

	if (filter)
		trace_destroy_filter(filter);
}

static void usage(char *prog) {
	fprintf(stderr, "Usage: %s [-i interval] [-f filter] inputURI\n",
		prog);
}


int main(int argc, char *argv[])
{
	libtrace_t *trace = NULL;
	libtrace_packet_t *packet = NULL;
	libtrace_filter_t *filter = NULL;

	int opt;
	char *filterstring = NULL;

	/* Ensure we have at least one argument after the program name */
	if (argc < 2) {
		usage(argv[0]);
		return 1;
	}

	/* Using getopt to handle any command line flags that would set the
	 * reporting interval and a filter */
	while ((opt = getopt(argc, argv, "i:f:")) != EOF) {
		switch (opt) {
			case 'i':
				interval = atoi(optarg);
				break;
			case 'f':
				filterstring = optarg;
				break;
			default:
				usage(argv[0]);
				return 1;
		}
	}
	
	/* After processing the options, we still need an argument to define
	 * the input URI */
	if (optind + 1 > argc) {
		usage(argv[0]);
		return 1;
	}
	
	/* Create the filter if a filter string was provided */
	if (filterstring != NULL) {
		filter = trace_create_filter(filterstring);
		if (filter == NULL) {
			fprintf(stderr, "Failed to create filter (%s)\n",
				filterstring);
			libtrace_cleanup(trace, packet, filter);
			return 1;
		}
	}

	/* Creating and initialising a packet structure to store the packets
	 * that we're going to read from the trace */
	packet = trace_create_packet();

	if (packet == NULL) {
		/* Unfortunately, trace_create_packet doesn't use the libtrace
		 * error system. This is because libtrace errors are associated
		 * with the trace structure, not the packet. In our case, we
		 * haven't even created a trace at this point so we can't 
		 * really expect libtrace to set an error on it for us, can
		 * we?
		 */
		perror("Creating libtrace packet");
		libtrace_cleanup(trace, packet, filter);
		return 1;
	}

	/* Opening and starting the input trace. Note that unlike the other
	 * examples, we can't just use argv[1] as we may have seen command
	 * line options. Instead we should use optind which will be set to
	 * the index of the first non-getopt argument  */
	trace = trace_create(argv[optind]);

	if (trace_is_err(trace)) {
		trace_perror(trace,"Opening trace file");
		libtrace_cleanup(trace, packet, filter);
		return 1;
	}

	/* Apply a filter, if one was created */
	if (filter != NULL) {
		if (trace_config(trace, TRACE_OPTION_FILTER, filter) == -1) {
			trace_perror(trace, "Configuring filter");
			libtrace_cleanup(trace, packet, filter);
			return 1;
		}
	}

	if (trace_start(trace) == -1) {
		trace_perror(trace,"Starting trace");
		libtrace_cleanup(trace, packet, filter);
		return 1;
	}

	/* This loop will read packets from the trace until either EOF is
	 * reached or an error occurs (hopefully the former!)
	 *
	 * Remember, EOF will return 0 so we only want to continue looping
	 * as long as the return value is greater than zero
	 */
	while (trace_read_packet(trace,packet)>0) {
		/* Call our per_packet function for every packet */
		per_packet(packet);
	}

	/* If the trace is in an error state, then we know that we fell out of
	 * the above loop because an error occurred rather than EOF being
	 * reached. Therefore, we should probably tell the user that something
	 * went wrong
	 */
	if (trace_is_err(trace)) {
		trace_perror(trace,"Reading packets");
		libtrace_cleanup(trace, packet, filter);
		return 1;
	}
	
	/* Print out the contents of the counters before exiting */
	print_stats();

	libtrace_cleanup(trace, packet, filter);

	return 0;
}