#include <linux/netfilter.h>
#include <netinet/ip.h>
#include <linux/tcp.h>
#include <linux/udp.h>
#include <libipq.h>
#include <stdio.h>

#define PACKET_LENGTH 2048

static char* __u32_to_str (__u32 ip)
{
   unsigned int tmp = 0;
   int i;
   char *str_ip;

   str_ip = (char *)malloc (sizeof(char) * 16);
   str_ip[0] = '\0';
   for (i=0; i<4; i++)
   {
      tmp = ip & 255;
      sprintf ((str_ip + strlen(str_ip)), "%d.", (unsigned int)tmp);
      ip >>= 8;
   }
   sprintf ((str_ip + (strlen(str_ip) - 1)), "%c", '\0');

   str_ip[15] = '\0';
   return ((char *)str_ip);
}


static void die(struct ipq_handle *h)
{
   ipq_perror("Error");
   ipq_destroy_handle(h);
   exit(1);
}

int main(int argc, char *argv[])
{
   int nStatus;
   unsigned char cbBuf[PACKET_LENGTH];
   struct ipq_handle* lpstHandle;
   struct ipq_packet_msg* lpstPacketMsg;
   struct iphdr* lpstIPHeader;
   struct tcphdr* lpstTCP;
   struct udphdr* lpstUDP;

   /* init and create queue */
   lpstHandle = ipq_create_handle(0, PF_INET);
   if (!lpstHandle) die(lpstHandle);

   /* mode: (here) copy entire packet to the user space and advise ip_queue */
   nStatus = ipq_set_mode(lpstHandle, IPQ_COPY_PACKET, PACKET_LENGTH);
   if (nStatus < 0) die(lpstHandle);

   while(1)
   {
     /* block for packet */
     nStatus = ipq_read(lpstHandle, cbBuf, PACKET_LENGTH, 0);
     if (nStatus < 0) die(lpstHandle);

     switch (ipq_message_type(cbBuf))
     {
        case NLMSG_ERROR:
        {
           fprintf(stderr, "%s: Received error message %d.\n", argv[0], ipq_get_msgerr(cbBuf));
           break;
        }
        case IPQM_PACKET:
        {
           /* get packet */
           lpstPacketMsg = ipq_get_packet(cbBuf);
           lpstIPHeader = ((struct iphdr*)lpstPacketMsg->payload);

           /* show packet info */
           if (lpstIPHeader && lpstIPHeader->protocol == IPPROTO_TCP)
           {
              lpstTCP = (struct tcphdr*)(lpstPacketMsg->payload + (lpstIPHeader->ihl << 2));
              if (lpstTCP)
              {
                 printf("[TCP] %s:%d -> %s:%d\n", __u32_to_str(lpstIPHeader->saddr),
                        lpstTCP->source, __u32_to_str(lpstIPHeader->daddr), lpstTCP->dest);
              }
           }
           else if (lpstIPHeader && lpstIPHeader->protocol == IPPROTO_UDP)
           {
              lpstUDP = (struct udphdr*)(lpstPacketMsg->payload + (lpstIPHeader->ihl << 2));
              if (lpstUDP)
              {
                 printf("[UDP] %s:%d -> %s:%d\n", __u32_to_str(lpstIPHeader->saddr),
                        lpstUDP->source, __u32_to_str(lpstIPHeader->daddr), lpstUDP->dest);
              }
           }

           /* verdict: (here) accept */
           nStatus = ipq_set_verdict(lpstHandle, lpstPacketMsg->packet_id, NF_ACCEPT, 0, NULL);
           if (nStatus < 0)  die(lpstHandle);
           break;
        }
        default:
        {
           fprintf(stderr, "%s: Unknown message type!\n", argv[0]);
           break;
        }
     }
   };

   ipq_destroy_handle(lpstHandle);
   return 0;
}