/* Compile with gcc -lnfnetlink -lnetfilter_queue  */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <linux/netfilter.h>		/* for NF_ACCEPT */
#include <arpa/inet.h>

#include <libnetfilter_queue/libnetfilter_queue.h>

#define  BUFSIZE  2048

struct in_addr foreign;
struct in_addr local;

struct queued_pckt {
	char *payload;
	int payload_len;
};

struct pseudohdr
{
	unsigned long ip_src ;
	unsigned long ip_dst ;
	unsigned char reserve ;
	unsigned char type ;
	unsigned short length;
} ;

unsigned short checksum(unsigned short *addr, unsigned int count) {
	/* Compute Internet Checksum for "count" bytes beginning at location "addr".
 	 * Algorithm is simple, using a 32-bit accumulator (sum),
	 * we add sequential 16-bit words to it, and at the end, fold back
	 * all the carry bits from the top 16 bits into the lower 16 bits.
 	*/
       register long sum = 0;
       unsigned short result;

        while (count > 1)  {
           /*  This is the inner loop */
               sum += * addr++;
               count -= 2;
       }
           /*  Add left-over byte, if any */
       if (count == 1)
       {
       		result = 0;	//make sure top half is zero
       		* (unsigned char *) (&result) = *(unsigned char *)addr;
       		sum += result;
	}
	
        /*
	 * Add back carry outs from top 16 bits to low 16 bits.
         * Fold 32-bit sum to 16 bits 
	*/
       
       sum = (sum >> 16) + (sum & 0xffff); 	/* add high-16 to low-16 */
       sum += (sum >> 16);			/* add carry */ 

       result = ~sum;			/* ones-complement, then truncate to 16 bits */	
       return (result);
}


unsigned short get_tcp_chksum (struct tcphdr *orig_tcphdr, struct iphdr *orig_iphdr )
{
	struct pseudohdr pseudoh ;

	unsigned int total_len = ntohs(orig_iphdr->tot_len);
	int tcpopt_len = (orig_tcphdr->doff * 4) - 20;
	int tcpdata_len = total_len - (orig_tcphdr->doff * 4) - (orig_iphdr->ihl * 4); 

	pseudoh.ip_src = orig_iphdr->saddr ;
	pseudoh.ip_dst = orig_iphdr->daddr ;
	pseudoh.reserve = 0 ; 	
	pseudoh.type = orig_iphdr->protocol ;  
	pseudoh.length = htons (sizeof (struct tcphdr) + tcpopt_len + tcpdata_len) ;

	int totaltcp_len = sizeof(struct pseudohdr) + sizeof(struct tcphdr) + tcpopt_len + tcpdata_len; 

	unsigned short *tcp = (unsigned short *)malloc (totaltcp_len);

	memcpy ((unsigned char *)tcp, &pseudoh, sizeof(struct pseudohdr));
	memcpy ((unsigned char *)tcp + sizeof(struct pseudohdr), (unsigned char *)orig_tcphdr, sizeof(struct tcphdr));
	if (tcpopt_len > 0)
		memcpy ((unsigned char *)tcp + sizeof(struct pseudohdr) + sizeof(struct tcphdr), (unsigned char *)orig_iphdr + (orig_iphdr->ihl * 4) + sizeof(struct tcphdr), tcpopt_len);
	
	if (tcpdata_len > 0)
		memcpy ((unsigned char *)tcp + sizeof(struct pseudohdr) + sizeof(struct tcphdr) + tcpopt_len, (unsigned char *)orig_tcphdr + (orig_tcphdr->doff * 4), tcpdata_len);

#if 0
	printf("pseudo length: %d\n",pseudoh.length);
        printf("tcp hdr length: %d\n",orig_tcphdr->doff*4);
        printf("tcp hdr struct length: %d\n",sizeof(struct tcphdr));
        printf("tcphdr->doff = %d, tcp opt length: %d\n",orig_tcphdr->doff,tcpopt_len);
        printf("tcp total+psuedo length: %d\n",totaltcp_len);

        fflush(stdout);

        printf("tcp data len: %d, data start %u\n", tcpdata_len,orig_tcphdr + (orig_tcphdr->doff*4));
#endif
	
	return (checksum (tcp, totaltcp_len)) ;
}

static void filter(
	unsigned char *packet, unsigned int payload_len)
{
	struct iphdr *iphdr;
	struct tcphdr *tcphdr;

	printf ("in filter function\n");

	iphdr = (struct iphdr *)packet;
	/* check need some datas */
	if (payload_len < sizeof(struct iphdr) + sizeof(struct tcphdr)) {
		return;
	}
	/* check IP version */
	if (iphdr->protocol == IPPROTO_TCP)
	{
		tcphdr = (struct tcphdr *)(((u_int32_t *)packet) + 4 * iphdr->ihl);

		if (iphdr->daddr == foreign.s_addr)
		{
			printf ("packet DEST addr = %s\n",inet_ntoa(foreign));
			fprintf (stderr, "changing pkt's DEST addr from FOREIGN to LOCAL\n");
			iphdr->daddr = local.s_addr;
			tcphdr->check = 0 ; // checksum will be calculated later
			iphdr->check = checksum(
				(unsigned short *)iphdr,
				sizeof(struct iphdr));
			tcphdr->check = get_tcp_chksum(
				tcphdr,
				iphdr);
		}
	}

}

static int cb(struct nfq_q_handle *qh, struct nfgenmsg *nfmsg,
	      struct nfq_data *nfa, void *data)
{
	int id = 0;
	struct nfqnl_msg_packet_hdr *ph;
	struct queued_pckt q_pckt;
	u_int32_t mark,ifi; 
	int ret;
	char *payload;
	
	printf("entering callback\n");
	ph = nfq_get_msg_packet_hdr(nfa);
	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(nfa);
	if (mark)
		printf("mark=%u ", mark);

	ifi = nfq_get_indev(nfa);
	if (ifi)
		printf("indev=%u ", ifi);

	ifi = nfq_get_outdev(nfa);
	if (ifi)
		printf("outdev=%u ", ifi);

	q_pckt.payload_len = nfq_get_payload(nfa, &(q_pckt.payload));
	if (q_pckt.payload_len >= 0)
	{
		printf("payload_len=%d ", q_pckt.payload_len);
		fputc('\n', stdout);
		filter((unsigned char *)q_pckt.payload, q_pckt.payload_len);
	}
	
	printf("setting verdict of packet id %d\n",id);
	return nfq_set_verdict(qh, id, NF_ACCEPT, q_pckt.payload_len, q_pckt.payload);
}

int main(int argc, char **argv)
{
	struct nfq_handle *h;
	struct nfq_q_handle *qh;
	struct nfnl_handle *nh;
	int fd;
	int rv;
	unsigned char buf[BUFSIZE];

	if (argc == 1)
	{
		inet_aton("10.102.130.222", &(foreign));
		inet_aton("10.102.130.105", &(local));
	} else if (argc == 3)
	{
		inet_aton(argv[1], &(foreign));
		inet_aton(argv[2], &(local));
	}
	else
	{
		printf("Usage: argv[0] [foreign_addr local_addr]\n");
		return 0;
	}

	printf("opening library handle\n");
	h = nfq_open();
	if (!h) {
		fprintf(stderr, "error during nfq_open()\n");
		return 0;
	}

	printf("unbinding existing nf_queue handler for AF_INET (if any)\n");
	if (nfq_unbind_pf(h, AF_INET) < 0) {
		fprintf(stderr, "error during nfq_unbind_pf()\n");
		exit(1);
	}

	printf("binding nfnetlink_queue as nf_queue handler for AF_INET\n");
	if (nfq_bind_pf(h, AF_INET) < 0) {
		fprintf(stderr, "error during nfq_bind_pf()\n");
		exit(1);
	}

	printf("binding this socket to queue '0'\n");
	qh = nfq_create_queue(h,  0, &cb, NULL);
	if (!qh) {
		fprintf(stderr, "error during nfq_create_queue()\n");
		exit(1);
	}

	printf("setting copy_packet mode\n");
	if (nfq_set_mode(qh, NFQNL_COPY_PACKET, BUFSIZE) < 0) {
		fprintf(stderr, "can't set packet_copy mode\n");
		exit(1);
	}

	nh = nfq_nfnlh(h);
	fd = nfnl_fd(nh);

	while ((rv = recv(fd, buf, BUFSIZE, 0)) && rv >= 0) {
		printf("pkt received\n");
		nfq_handle_packet(h, buf, rv);
		printf("pkt handled\n");
	}

	printf("unbinding from queue 0\n");
	nfq_destroy_queue(qh);

	printf("closing library handle\n");
	nfq_close(h);

	exit(0);
}