#include <stdio.h>
#include <netdb.h>
#include <string.h>
#include <stdlib.h>
#include <getopt.h>
#include <iptables.h>
#include <linux/netfilter_ipv4/ipt_ipaddr.h>

/* 
 * iptables -m ipt_ipaddr -h 
 */
static void help(void)
{
   printf (
            "IPADDR v%s options:\n"
            "[!] --ipsrc <ip>\t\t The incoming ip addr matches.\n"
            "[!] --ipdst <ip>\t\t The outgoing ip addr matches.\n"
            "\n", IPTABLES_VERSION
         );
}

static struct option opts[] = {
   { .name = "ipsrc",   .has_arg = 1,   .flag = 0,   .val = '1' },
   { .name = "ipdst",   .has_arg = 1,   .flag = 0,   .val = '2' },
   { .name = 0 }
};

static void init(struct ipt_entry_match *m, unsigned int *nfcache)
{
   /* Can't cache this */
   *nfcache |= NFC_UNKNOWN;
}

static void parse_ipaddr(const char *arg, u_int32_t *ipaddr)
{
   struct in_addr *ip;

   ip = dotted_to_addr(arg);
   if (!ip)
      exit_error(PARAMETER_PROBLEM, "ipt_ipaddr: Bad IP address `%s'\n", arg);

   *ipaddr = ip->s_addr;
}

static int parse(int c, char **argv, int invert, unsigned int *flags,
                 const struct ipt_entry *entry,
                 unsigned int *nfcache,
                 struct ipt_entry_match **match)
{
   struct ipt_ipaddr_info *info = (struct ipt_ipaddr_info *)(*match)->data;

   switch(c) {
      case '1':
         if (*flags & IPADDR_SRC)
            exit_error(PARAMETER_PROBLEM, "ipt_ipaddr: Only use --ipsrc once!");
         *flags |= IPADDR_SRC;

         info->flags |= IPADDR_SRC;
         if (invert)
            info->flags |= IPADDR_SRC_INV;

         parse_ipaddr(argv[optind-1], &info->ipaddr.src);
         break;

      case '2':
         if (*flags & IPADDR_DST)
            exit_error(PARAMETER_PROBLEM, "ipt_ipaddr: Only use --ipdst once!");
         *flags |= IPADDR_DST;

         info->flags |= IPADDR_DST;
         if (invert)
            info->flags |= IPADDR_DST_INV;

         parse_ipaddr(argv[optind-1], &info->ipaddr.dst);
         break;

      default:
         return 0;
   }

   return 1;
}

static void final_check(unsigned int flags)
{
   if (!flags)
      exit_error(PARAMETER_PROBLEM, "ipt_ipaddr: Invalid parameters.");
}

/* 
 * Can't call it ipaddr, because
 * the main structure is already
 * called like that.
 */
static void print_ipaddr(u_int32_t *ip_addr)
{
   printf(  "%u.%u.%u.%u ", 
            ((unsigned char *)ip_addr)[0],
            ((unsigned char *)ip_addr)[1],
            ((unsigned char *)ip_addr)[2],
            ((unsigned char *)ip_addr)[3]
         );
}

static void print(const struct ipt_ip *ip,
                  const struct ipt_entry_match *match,
                  int numeric)
{
   const struct ipt_ipaddr_info *info = (const struct ipt_ipaddr_info *)match->data;

   if (info->flags & IPADDR_SRC) {
         printf("src IP ");
      if (info->flags & IPADDR_SRC_INV)
         printf("! ");
      print_ipaddr((u_int32_t *)&info->ipaddr.src);
   }

   if (info->flags & IPADDR_DST) {
      printf("dst IP ");
      if (info->flags & IPADDR_DST_INV)
         printf("! ");
      print_ipaddr((u_int32_t *)&info->ipaddr.dst);
   }
}

/* 
 * iptables-saves
 */
static void save(const struct ipt_ip *ip, const struct ipt_entry_match *match)
{
   const struct ipt_ipaddr_info *info = (const struct ipt_ipaddr_info *)match->data;

   if (info->flags & IPADDR_SRC) {
      if (info->flags & IPADDR_SRC_INV)
         printf("! ");
      printf("--ipsrc ");
      print_ipaddr((u_int32_t *)&info->ipaddr.src);
   }

   if (info->flags & IPADDR_DST) {
      if (info->flags & IPADDR_DST_INV)
         printf("! ");
      printf("--ipdst ");
      print_ipaddr((u_int32_t *)&info->ipaddr.dst);
   }
}

static struct iptables_match ipaddr
= {
    .name            = "ipaddr",
    .version         = IPTABLES_VERSION,
    .size            = IPT_ALIGN(sizeof(struct ipt_ipaddr_info)),
    .userspacesize   = IPT_ALIGN(sizeof(struct ipt_ipaddr_info)),
    .help            = &help,
    .init            = &init,
    .parse           = &parse,
    .final_check     = &final_check,
    .print           = &print,
    .save            = &save,
    .extra_opts      = opts
};

void _init(void)
{
   register_match(&ipaddr);
}