#include <linux/netfilter.h>
#include <libipq.h>
#include <stdio.h>
#include <syslog.h>
#include <GeoIP.h>
#include <netinet/ip.h>

#define BUFSIZE 2048

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

char *numtochar(unsigned long ipnum, char *mychar)
{
  unsigned long a1,a2,a3,a4; 

  a1 = ipnum >> 24;
  a2 = (ipnum << 8) >> 24;
  a3 = (ipnum << 16) >> 24;
  a4 = (ipnum << 24) >> 24;
  sprintf(mychar, "%u.%u.%u.%u", a1, a2, a3, a4);
  return(mychar);
}

void usage() {
  fprintf(stderr,"Usage: chinaspam <countrycode1> [countrycode2] ...\n");
  fprintf(stderr,"Example: chinaspam CN TW HK\n");
  fprintf(stderr,"Check ISO 3166 countrycodes at http://www.maxmind.com/app/iso3166\n");
}

int main(int argc, char **argv)
{
  int status, countryid, match, i;
  unsigned char buf[BUFSIZE], *payload;
  ipq_packet_msg_t *m;
  struct ipq_handle *h;
  GeoIP * gi;
  char *countrycode, *country, char_ip[15];
  unsigned long ipnum;

  if (argc == 1) {
    usage();
    exit(1);
  }

  gi = GeoIP_new(GEOIP_STANDARD);

  // for iptables version 1.2.5 and newer :
  h = ipq_create_handle(0, PF_INET);
  // for iptables version 1.2.4 and older :
  // h = ipq_create_handle(0);

  if (!h)
    die(h);

  status = ipq_set_mode(h, IPQ_COPY_PACKET, BUFSIZE);
  if (status < 0)
    die(h);

  do{
    status = ipq_read(h, buf, BUFSIZE, 0);
    if (status < 0)
      die(h);

    switch (ipq_message_type(buf)) {
      case NLMSG_ERROR:
        fprintf(stderr, "Received error message %d\n", ipq_get_msgerr(buf));
        break;

      case IPQM_PACKET: {
        m = ipq_get_packet(buf);
	payload = m->payload;
	ipnum = htonl(((struct iphdr *)payload)->saddr);

	countryid = GeoIP_id_by_ipnum(gi, ipnum);
	countrycode = (countryid > 0) ? (char *) GeoIP_country_code[countryid] : NULL;

        match = 0;
        for (i=1 ; ((i<argc) && (match==0)) ; i++) {
          if ( strcmp(countrycode, argv[i]) == 0 )
            match=1;
        }

	if ( match != 0 ) {
          status = ipq_set_verdict(h, m->packet_id, NF_DROP, 0, NULL);
	  country = (countryid > 0) ? (char *) GeoIP_country_name[countryid] : NULL;
	  syslog(LOG_DAEMON|LOG_INFO, "Dropped a packet from %s because it came from non-authorized country : %s", numtochar(ipnum, char_ip), country);
          if (status < 0)
            die(h);
          break;
        }
        else {
          status = ipq_set_verdict(h, m->packet_id, NF_ACCEPT, 0, NULL);
          if (status < 0)
            die(h);
          break;
        }
      }
      default:
        fprintf(stderr, "Unknown message type!\n");
        break;
    }
  } while (1);
  ipq_destroy_handle(h);
  return 0;
}