#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <netinet/udp.h>
#include <sys/ioctl.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/time.h>
#include <arpa/inet.h>

#include <string.h>
#include <signal.h>

#include <map>
#include <string>
#include <sstream>

using namespace std;


#ifdef  __cplusplus
extern "C" {
#endif

#include <linux/netfilter.h>		/* for NF_ACCEPT */
#include <libnfnetlink/libnfnetlink.h>
#include <libnetfilter_queue/libnetfilter_queue.h>
#include <libnetfilter_queue/linux_nfnetlink_queue.h>

#ifdef  __cplusplus
}
#endif


#define QNUM 40
#define QSTART 1000

int TERM=0;

void sig_term(int sig) {
 TERM=1;
}

struct info {
 int qnum;
 struct nfq_q_handle *qh;
};

map<int,info> queues; //queus[qnum]=info

//GLOBAlS HERE
struct nfq_handle *h;
struct nfq_q_handle *qh;
struct nfnl_handle *nh;
int nfqfd;

int rv;
char buf[4096];
//=========================

template <class T>
std::string stringify(T x) {
 std::stringstream o;
 o << x;
 return o.str();
}

#include <sys/resource.h>

int force_core(unsigned long core_size_cur, unsigned long core_size_max) {
 struct  rlimit rlim;

 rlim.rlim_cur=core_size_cur;
 rlim.rlim_max=core_size_max;
 return setrlimit(RLIMIT_CORE,&rlim);
}

//HANDLES PACKETS HERE
static u_int32_t handle_pkt (struct nfq_data *tb,struct info *spec,int &verdict) {
  int id = 0;
  struct nfqnl_msg_packet_hdr *ph;
  u_int32_t mark;
  int ret;
  char *data;
//=	
  struct iphdr  *ip;
  struct tcphdr *tcp;
  struct udphdr *udp;
  char saddr[20],daddr[20];
  int sport=0;
  int dport=0;
  int i;
//=	
  
  verdict=NF_ACCEPT;  
	
  ph = nfq_get_msg_packet_hdr(tb);
  if (ph) {
	id = ntohl(ph->packet_id);
	printf("hw_protocol=0x%04x hook=%u id=%u ",
	ntohs(ph->hw_protocol), ph->hook, id);	 			
  }
	
  mark = nfq_get_nfmark(tb);
  printf("mark=%u ", mark);
  ret = nfq_get_payload(tb, &data);

//======
  ip=(struct iphdr*) data;
  if (ip->protocol==6) {
    tcp=(struct tcphdr*) (data + (4 * ip->ihl));
    sport = htons(tcp->source);
    dport = htons(tcp->dest);
  } else if (ip->protocol==17)  {
     udp=(struct udphdr*) (data + (4 * ip->ihl));
     sport = htons(udp->source);
     dport = htons(udp->dest);
  }

  strcpy(saddr,inet_ntoa(*(struct in_addr*)&ip->saddr));
  strcpy(daddr,inet_ntoa(*(struct in_addr*)&ip->daddr));

  printf("%i src=%s:%u  dst=%s:%u size=%u proto=%u",
                 spec->qnum,saddr,sport,daddr,dport,htons(ip->tot_len),ip->protocol);

//=======
  fputc('\n', stdout);
  return id;
}
	

static int cb(struct nfq_q_handle *qh, struct nfgenmsg *nfmsg,
	      struct nfq_data *nfa, void *data) {
        int verdict=NF_DROP;	      
 	u_int32_t id = handle_pkt(nfa, (struct info *) data, verdict);
	return nfq_set_verdict(qh, id, verdict, 0, NULL);
}


//INIT NFQ
void init_nfq() {
 h=0;
 
 printf("opening library handle\n");
 h = nfq_open();
 if (!h) throw "error during nfq_open()";

 printf("unbinding existing nf_queue handler for AF_INET (if any)\n");
 if (nfq_unbind_pf(h, AF_INET) < 0) throw "error during nfq_unbind_pf()";

 printf("binding nfnetlink_queue as nf_queue handler for AF_INET\n");
 if (nfq_bind_pf(h, AF_INET) < 0) throw "error during nfq_bind_pf()";

 nh = nfq_nfnlh(h);
 nfqfd = nfnl_fd(nh);
    
 if (nfqfd>0) {
   fcntl(nfqfd,F_SETFL,O_NONBLOCK);
 } else throw "fail to set nfq nfnl fd";
}


//CHECK IF THERE ANY PACKETS IN QUEUE
int check_nfq(int timeout) {
 fd_set readfds;
 struct timeval tv;
 int res;

 if (!queues.size()) return -100;

 FD_ZERO(&readfds);
 FD_SET(nfqfd,&readfds);

 if (timeout<0) {
   res=select(nfqfd+1,&readfds,NULL,NULL,NULL);
 } else {
   tv.tv_sec=0; //timeout;
   tv.tv_usec=timeout;
   res=select(nfqfd+1,&readfds,NULL,NULL,&tv);
 }

 if (res>0) {
  res=0;
   int rv;
   char buf[4096];
   if (FD_ISSET(nfqfd,&readfds)) {
    res=1;
    rv = recv(nfqfd, buf, sizeof(buf), 0);
    if (errno <0) throw strerror(errno);
    nfq_handle_packet(h, buf, rv);
   }
 }
 
 return res;
 
}



void create_queue(int qnum) {
  struct info spec;
  struct info *pspec;
  map<int,info>::iterator ii; 
  
  if ((ii=queues.find(qnum))!=queues.end()) {
   printf("queue '%i' already queued\n",ii->first);
   return;
  }

  spec.qnum=qnum;
  queues[spec.qnum]=spec;

  pspec=&queues[spec.qnum];
  printf("binding this socket to queue '%i'\n",spec.qnum);
  pspec->qh = nfq_create_queue(h,  spec.qnum, &cb, pspec);
  if (!pspec->qh) throw "error during nfq_create_queue()";
  printf("setting copy_packet mode\n");
  if (nfq_set_mode(pspec->qh, NFQNL_COPY_PACKET, 0xffff) < 0) {
    fprintf(stderr,"can't set packet_copy mode (%s)\n",strerror(errno));
    queues.erase(queues.find(spec.qnum));
  }  
}

void delete_queue(int qnum) {
 map<int,info>::iterator ii; 
 if ((ii=queues.find(qnum))!=queues.end()) {
   printf("destroing queue '%i'\n",ii->first);
   nfq_destroy_queue(ii->second.qh);
 }
}

void enable_access(int qnum,int order, bool enable=true) {
 char what;
 
 if (enable) what='A';
 else what='D';
 
 string cmd=string("iptables -") + what +" OUTPUT -d 192.168.1." + stringify(order)+  " -j NFQUEUE --queue-num "+stringify(qnum);
 puts(cmd.c_str());
 system(cmd.c_str());
 
}

//RUN TEST

#define TIMEOUT 10000

void run_test() {
 int i,n;
 map<int,info>::iterator ii; 
 n=0;
 
 force_core(-1,-1); 
 srand(time(0));
 
 for (i=0; i<QNUM; i++) { //init queues
  create_queue(QSTART+i);    
  enable_access(QSTART+i,i);
 }
 
 printf("Initialized %i queues\n",queues.size());
 
 time_t now=time(0);
 time_t last=0;
 while(!TERM) {
 //creating and destroying queues 
  check_nfq(TIMEOUT);
  now=time(0);
  if (now<last+2) continue;
  last=now;
  int x=(int) (40.0+160.0/RAND_MAX*rand()); //random 40..200
//  printf("%i\n",x);
//  continue;
  if (x & 1) { //try to create if x is odd
   create_queue(x+QSTART);
   enable_access(QSTART+x,x);
   continue;
  }
  if (queues.size()<QNUM) continue;
  delete_queue(x+QSTART); //try to create if x is even
  enable_access(QSTART+x,x,false);
 }
 
 for (ii=queues.begin();ii!=queues.end(); ++ii) { //destroy queues
  delete_queue(ii->first);
  enable_access(ii->first,ii->first-QSTART,false);  
 }
 
}

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

 signal(SIGTERM,sig_term);
 signal(SIGINT,sig_term);
 
 try { 
  init_nfq();
 } 
 catch(const char *msg) {
  fprintf(stderr,"Fail to init nfq: %s\n",msg);  
  exit(-1);
 } 

 try { 
  run_test();
 } 
 catch(const char *msg) {
  fprintf(stderr,"Failure during test running: %s\n",msg);  
  exit(-1);
 } 
 
 
 signal(SIGTERM,SIG_DFL);
 signal(SIGINT,SIG_DFL);


#ifdef INSANE
	/* normally, applications SHOULD NOT issue this command, since
	 * it detaches other programs/sockets from AF_INET, too ! */
	printf("unbinding from AF_INET\n");
	nfq_unbind_pf(h, AF_INET);
#endif
 printf("closing library handle\n");
 nfq_close(h);
 
 return 0;
}