#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <inttypes.h>
#include <string.h>
#include <signal.h>
#include <pcap.h>
#include <time.h>

#include <linux/limits.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/xattr.h>
#include <linux/capability.h>


#include "protocol_header.h"
#include "inode.h"
#include "inodelist.h"
#include "cnx_inode.h"
#include "mypcap.h"

#define APP_NAME "cleanpcap"
#define APP_VERSION "0.0.1"
#define APP_DATE "2014-12-04"
#define APP_DESCRIPTION "Small tool to dump pcaps split by pid"

#define MAX_CAPLEN 65535

#ifndef MAX_SIZE
	#define MAX_SIZE 256
#endif //MAX_SIZE

#ifndef MAX_FILTER_SIZE
	#define MAX_FILTER_SIZE 256
#endif //MAX_FILTER_SIZE

static char st_device[MAX_SIZE]="eth0";//By default launch on eth0
static char st_filter[MAX_FILTER_SIZE]="ip";//By filter only on IP datagram
static char st_file_prefix[MAX_SIZE]="";//Will be time of day
static uint64_t st_packet_num = 0;
//~ static pcap_t *st_hdl_out_pcap = NULL;
//~ static pcap_dumper_t *st_hdl_out_pcapdumper = NULL;
static FILE *st_hdl_out_mypcap = NULL;
static pcap_t *st_hdl_capture = NULL;

static void handle_sigint(int signo)
{
	(void)signo;
	if (st_hdl_capture != NULL) {
		pcap_breakloop(st_hdl_capture);
	}
}

static void close_output_files(void)
{
	for (inodelist_t *node = inodelist_get_head(); node != NULL; node = node->next) {
		if (node->data.mydumper != NULL) {
			mypcap_flush(node->data.mydumper);
			mypcap_close(node->data.mydumper);
			node->data.mydumper = NULL;
		}
	}
	if (st_hdl_out_mypcap != NULL) {
		mypcap_flush(st_hdl_out_mypcap);
		mypcap_close(st_hdl_out_mypcap);
		st_hdl_out_mypcap = NULL;
	}
}

static void print_capture_progress(void)
{
	printf("\r\033[KPackets captured: %" PRIu64, st_packet_num);
	fflush(stdout);
}

void usage()
{	
	printf( "%s (%s) %s\nDescription :\n\t%s\nUsage :\n\t%s <incap>\n", APP_NAME, APP_VERSION, APP_DATE, APP_DESCRIPTION, APP_NAME );
}

static char errbuff[PCAP_ERRBUF_SIZE]="";


void pcap_loop_cb(u_char *useless,const struct pcap_pkthdr* header,const u_char*
        packet)
{
	(void)useless;
	if (header == NULL || packet == NULL) {
		return;
	}
	++st_packet_num;
	uint64_t report_interval = st_packet_num <= 10 ? 1 :
		(st_packet_num <= 100 ? 10 : (st_packet_num <= 1000 ? 100 : 1000));
	if (st_packet_num % report_interval == 0) {
		print_capture_progress();
	}

	//Inode <-> Process Update
	cnx_inode_refresh();
	reread_mapping();
	
	if( NULL != st_hdl_out_mypcap ) {
		mypcap_dump(st_hdl_out_mypcap,
			(const struct pcap_pkthdr*)header,
			(uint8_t*)packet);
		if(0 == (st_packet_num % 10))
			mypcap_flush(st_hdl_out_mypcap);
	}
	if (header->caplen < ETHER_HEADER_SIZE) {
		return;
	}
	const size_t packet_len = header->caplen;
	size_t ip_offset = ETHER_HEADER_SIZE;
	const struct sniff_ethernet *eth_header = (const struct sniff_ethernet *)packet;
	if (ntohs(eth_header->ether_type) == 0x8100 || ntohs(eth_header->ether_type) == 0x88a8) {
		if (packet_len < ETHER_HEADER_SIZE + 4) {
			return;
		}
		ip_offset += 4;
	}
	if (packet_len < ip_offset + IP4_HEADER_SIZE) {
		return;
	}
	const sniff_ip_t *ip_header = (const sniff_ip_t *)(packet + ip_offset);
	size_t ip_header_len = IP_HL(ip_header) * 4U;
	if (IP_V(ip_header) != 4 || ip_header_len < IP4_HEADER_SIZE ||
		ip_header->ip_p != IPPROTO_UDP || packet_len < ip_offset + ip_header_len + UDP_HEADER_SIZE) {
		return;
	}
	const sniff_udp_t *udp_header = (const sniff_udp_t *)(packet + ip_offset + ip_header_len);

	char str_ip_src[STR_IP4_SIZE+1]="";
	char str_ip_dst[STR_IP4_SIZE+1]="";
	strncpy(str_ip_src,inet_ntoa(ip_header->ip_src),STR_IP4_SIZE);
	strncpy(str_ip_dst,inet_ntoa(ip_header->ip_dst),STR_IP4_SIZE);
	
	char hashkey[HASHKEYSIZE];
	snprintf(hashkey, HASHKEYSIZE * sizeof(char), "%s:%d-%s:%d",
			str_ip_src,
			ntohs(udp_header->uh_sport),
			str_ip_dst,
			ntohs(udp_header->uh_dport));
	
	//~ inodelist_debug_display_forward();
	
	//Search packet hashkey in the inodelist
	inodelist_t *result = inodelist_find_hashkey(hashkey);
	if( NULL != result ) {
		//Dump into its own file
		
		//~ printf("\n\t\t%s FOR\n", result->data.cmdline);
		printf("\r\033[K#%" PRIu64 " %s\n", st_packet_num, result->data.cmdline);
		print_capture_progress();
		//~ if( NULL != result->data.dumper ) {
			//~ pcap_dump((uint8_t*)result->data.dumper,
				//~ (const struct pcap_pkthdr*)header,
				//~ (uint8_t*)packet);
			//~ pcap_dump_flush(result->data.dumper);
		//~ }
		if( NULL != result->data.mydumper ) {
			mypcap_dump(result->data.mydumper,
				(const struct pcap_pkthdr*)header,
				(uint8_t*)packet);
			mypcap_flush(result->data.mydumper);
		}
	} else {
		snprintf(hashkey, HASHKEYSIZE * sizeof(char), "%s:%d-%s:%d",
			str_ip_src,
			ntohs(udp_header->uh_sport),
			str_ip_dst,
			ntohs(udp_header->uh_dport));
		result = inodelist_find_hashkey(hashkey);
		if( NULL != result ) {
			printf("\r\033[K#%" PRIu64 " %s reverse\n", st_packet_num, result->data.cmdline);
			print_capture_progress();
			if( NULL != result->data.mydumper ) {
				mypcap_dump(result->data.mydumper,
					(const struct pcap_pkthdr*)header,
					(uint8_t*)packet);
				mypcap_flush(result->data.mydumper);
			}
		} 
	}
}

int32_t main(int argc, char **argv )
{
	//~ inode_t inode1={1,101,"inode1"};
	//~ inode_t inode2={2,102,"inode2"};
	//~ inode_t inode3={3,103,"inode3"};
	//~ inode_t inode4={4,103,"inode4"};
	//~ inodelist_insert_head(inode1);
	//~ inodelist_insert_head(inode2);
	//~ inodelist_insert_head(inode3);
	//~ inodelist_insert_head(inode4);
	//~ inodelist_debug_display_forward();
	//~ inodelist_debug_display_backward();
	//~ inodelist_remove_at(1);
	//~ int idx = inodelist_get_index_from_value(3);
	//~ printf("idx : %d\n", idx);
	//~ inodelist_t *iltmp = inodelist_get_inode_from_value(3);
	//~ inode_debug_print(&iltmp->data);
	//~ inodelist_debug_display_forward();

	//Some checks
	if (geteuid() != 0) {
		char exe_path[PATH_MAX];
		ssize_t len;
		unsigned int caps[5] = {0,0,0,0,0};
		if ((len = readlink("/proc/self/exe", exe_path, PATH_MAX)) == -1) {
			fprintf(stderr, "ERROR: Cannot read self :(");
			exit(1);
		}
		exe_path[len] = '\0';
		getxattr(exe_path, "security.capability", (char *)caps, sizeof(caps));
		if( (((caps[1] >> CAP_NET_ADMIN) & 1) != 1) ||
			(((caps[1] >> CAP_NET_RAW) & 1) != 1) ) {
			fprintf(stderr, "ERROR: Must be root or tweak capabilities (cap_net_admin,cap_net_raw)\n");
			exit(2);
		}
	}
	
	if( argc > 1 ) {
		if (strlen(argv[1]) >= sizeof(st_device)) {
			fprintf(stderr, "ERROR: Interface name is too long\n");
			return 2;
		}
		strcpy(st_device, argv[1]);
	}
	if( argc > 2 ) {
		if (strlen(argv[2]) >= sizeof(st_filter)) {
			fprintf(stderr, "ERROR: Filter is too long\n");
			return 2;
		}
		strcpy(st_filter, argv[2]);
	}
	
	struct bpf_program fp;		/* The compiled filter expression */
	bpf_u_int32 mask;		/* The netmask of our sniffing device */
	bpf_u_int32 net;		/* The IP of our sniffing device */
	
	if (pcap_lookupnet(st_device, &net, &mask, errbuff) == -1) {
		fprintf(stderr, "WARNING: Can't get netmask for device %s\n", st_device);
		net = 0;
		mask = 0;
	}
	
	st_hdl_capture = pcap_open_live(st_device, MAX_CAPLEN, 0, 1000, errbuff);
	if (st_hdl_capture == NULL) {
		fprintf(stderr, "ERROR: Cannot open %s for capture\n", st_device);
		exit(3);
	}
	struct sigaction sigint_action;
	memset(&sigint_action, 0, sizeof(sigint_action));
	sigint_action.sa_handler = handle_sigint;
	sigemptyset(&sigint_action.sa_mask);
	if (sigaction(SIGINT, &sigint_action, NULL) == -1) {
		perror("ERROR: Cannot install SIGINT handler");
		pcap_close(st_hdl_capture);
		return 8;
	}

	if (pcap_compile(st_hdl_capture, &fp, st_filter, 0, net) == -1) {
		fprintf(stderr, "ERROR: Couldn't parse filter [%s]: %s\n", st_filter, pcap_geterr(st_hdl_capture));
		exit(4);
	}
	
	if (pcap_setfilter(st_hdl_capture, &fp) == -1) {
		fprintf(stderr, "ERROR: Couldn't install filter [%s]: %s\n", st_filter, pcap_geterr(st_hdl_capture));
		exit(5);
	}
	pcap_freecode(&fp);
	
	//Set file output prefix
	time_t     now;
    struct tm  ts;
    char       buf[80];
    // Get current time
    time(&now);
    // Format time, "ddd yyyy-mm-dd hh:mm:ss zzz"
    ts = *localtime(&now);
    strftime(buf, sizeof(buf), "pcap-%Y-%m-%d_%H%M%S_%Z", &ts);
    strcpy(st_file_prefix,buf);
	set_file_prefix(st_file_prefix);
	
    //Prepare complete dump file
    char complete_dump_path[MAX_SIZE]="";
    sprintf(complete_dump_path,"%s-complete.pcap", st_file_prefix);
	st_hdl_out_mypcap=mypcap_open(complete_dump_path, "w+");
	if (st_hdl_out_mypcap == NULL) {
		fprintf(stderr, "ERROR: Cannot create %s\n", complete_dump_path);
		pcap_close(st_hdl_capture);
		return 7;
	}
	
	//Inode <-> Process Init
	cnx_inode_refresh();//Create a list of all inode handling socket
	reread_mapping();//Map pid to inodes found
	
	
		
	printf("listening on %s, capture size %d bytes\n",st_device,MAX_CAPLEN);
	printf("filtering on [%s]\n",st_filter);
	printf("complete dump file : %s\n",complete_dump_path);
	if (pcap_datalink(st_hdl_capture) != DLT_EN10MB) {
		fprintf(stderr, "ERROR: No Ethernet headers - not supported\n");
		close_output_files();
		pcap_close(st_hdl_capture);
		exit(6);
	} else {
		printf("link-type EN10MB (Ethernet)\n");
	}
	int loop_result = pcap_loop(st_hdl_capture, 0, pcap_loop_cb, NULL);
	if (loop_result == -1) {
		fprintf(stderr, "ERROR: Capture failed: %s\n", pcap_geterr(st_hdl_capture));
	}
	printf("\n");
	close_output_files();
	pcap_close(st_hdl_capture);
	
#ifdef _DEBUG
	printf("\n");
#endif //_DEBUG
	return 0;
}