#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.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 int 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;

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)
{
	//Inode <-> Process Update
	cnx_inode_refresh();
	reread_mapping();
	
	++st_packet_num;//update packet counter
	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);
	}
	
	struct sniff_ethernet *eth_header = (struct sniff_ethernet *)packet;
	uint16_t vlan_shift = 0;
	if( eth_header->ether_type == 0x0081 ) {
		vlan_shift = 4;
		packet = packet + vlan_shift;
	}
		
	sniff_ip_t *ip_header = (sniff_ip_t*)(packet+ETHER_HEADER_SIZE);
	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);
	
	sniff_udp_t *udp_header = (sniff_udp_t*)(packet+ETHER_HEADER_SIZE+IP4_HEADER_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("#%i %s\n", st_packet_num, result->data.cmdline);
		//~ 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("#%i %s reverse\n", st_packet_num, result->data.cmdline);
			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 ) {
		strcpy(st_device, argv[1]);
	}
	if( argc > 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;
	}
	
	pcap_t *pcap = pcap_open_live(st_device, MAX_CAPLEN, 0, 1000, errbuff);
	if (pcap == NULL) {
		fprintf(stderr, "ERROR: Cannot open %s for capture\n", st_device);
		exit(3);
	}

	if (pcap_compile(pcap, &fp, st_filter, 0, net) == -1) {
		fprintf(stderr, "ERROR: Couldn't parse filter [%s]: %s\n", st_filter, pcap_geterr(pcap));
		exit(4);
	}
	
	if (pcap_setfilter(pcap, &fp) == -1) {
		fprintf(stderr, "ERROR: Couldn't install filter [%s]: %s\n", st_filter, pcap_geterr(pcap));
		exit(5);
	}
	
	//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+");
	
	//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(pcap) != DLT_EN10MB) {
		fprintf(stderr, "ERROR: No Ethernet headers - not supported\n");
		exit(6);
	} else {
		printf("link-type EN10MB (Ethernet)\n");
	}
	
	pcap_loop(pcap, 0, pcap_loop_cb, NULL );
	pcap_loop(pcap, 100, pcap_loop_cb, NULL );

	mypcap_close(st_hdl_out_mypcap);
	pcap_close( pcap );
	
#ifdef _DEBUG
	printf("\n");
#endif //_DEBUG
	return 0;
}