Python ext for iptables-save

iptables_save.c

#include <Python.h>
#include <libiptc/libiptc.h>
#include <xtables.h>
#include <stdio.h>
#include <netdb.h>
#include <string.h>


#define buffer_printf(...) {\
  char temp_buf[1024];\
  snprintf(temp_buf, sizeof(temp_buf), __VA_ARGS__);\
  size_t temp_buf_len = strlen(temp_buf), buffer_len = strlen(*pbuffer);\
  if (temp_buf_len > 0) {\
    size_t new_size = buffer_len + temp_buf_len + 1;\
    *pbuffer = realloc(*pbuffer, new_size);\
    strcat(*pbuffer, temp_buf);\
  }\
}

// {{{
// Taken from iptables project, commit 03a3f7cf1826f70dd1db070aa8ee11bd2abbccd9, and
// - do not use 'alias' because older xtables (prior to 983196ceb4d3bb7b6d3cf6da18bb6d5a5eafb347) doesn't have it
// - silently ignore errors, remove fprintf(stderr, ...)
// - replace printf to a customized one, buffer_printf, to capture the output

#define IP_PARTS_NATIVE(n)      \
  (unsigned int)((n)>>24)&0xFF,      \
  (unsigned int)((n)>>16)&0xFF,      \
  (unsigned int)((n)>>8)&0xFF,      \
  (unsigned int)((n)&0xFF)

#define IP_PARTS(n) IP_PARTS_NATIVE(ntohl(n))

static void print_proto(uint16_t proto, int invert, char **pbuffer)
{
  if (proto) {
    unsigned int i;
    const char *invertstr = invert ? " !" : "";

    const struct protoent *pent = getprotobynumber(proto);
    if (pent) {
      buffer_printf("%s -p %s", invertstr, pent->p_name);
      return;
    }

    for (i = 0; xtables_chain_protos[i].name != NULL; ++i)
      if (xtables_chain_protos[i].num == proto) {
        buffer_printf("%s -p %s",
               invertstr, xtables_chain_protos[i].name);
        return;
      }

    buffer_printf("%s -p %u", invertstr, proto);
  }
}

/* print a given ip including mask if neccessary */
static void print_ip(const char *prefix, uint32_t ip,
         uint32_t mask, int invert, char **pbuffer)
{
  uint32_t bits, hmask = ntohl(mask);
  int i;

  if (!mask && !ip && !invert)
    return;

  buffer_printf("%s %s %u.%u.%u.%u",
    invert ? " !" : "",
    prefix,
    IP_PARTS(ip));

  if (mask == 0xFFFFFFFFU) {
    buffer_printf("/32");
    return;
  }

  i    = 32;
  bits = 0xFFFFFFFEU;
  while (--i >= 0 && hmask != bits)
    bits <<= 1;
  if (i >= 0)
    buffer_printf("/%u", i)
  else
    buffer_printf("/%u.%u.%u.%u", IP_PARTS(mask));
}

/* This assumes that mask is contiguous, and byte-bounded. */
static void
print_iface(char letter, const char *iface, const unsigned char *mask,
      int invert, char **pbuffer)
{
  unsigned int i;

  if (mask[0] == 0)
    return;

  buffer_printf("%s -%c ", invert ? " !" : "", letter);

  for (i = 0; i < IFNAMSIZ; i++) {
    if (mask[i] != 0) {
      if (iface[i] != '\0')
        buffer_printf("%c", iface[i]);
    } else {
      /* we can access iface[i-1] here, because
       * a few lines above we make sure that mask[0] != 0 */
      if (iface[i-1] != '\0')
        buffer_printf("+");
      break;
    }
  }
}

static int print_match_save(const struct xt_entry_match *e,
      const struct ipt_ip *ip, char **pbuffer)
{
  const struct xtables_match *match =
    xtables_find_match(e->u.user.name, XTF_TRY_LOAD, NULL);

  if (match) {
    buffer_printf(" -m %s", e->u.user.name);

    /* some matches don't provide a save function */
    if (match->save)
      match->save(ip, e);
  } else {
    if (e->u.match_size) {
      fprintf(stderr,
        "Can't find library for match `%s'\n",
        e->u.user.name);
      return -1;
    }
  }
  return 0;
}

/* We want this to be readable, so only print out neccessary fields.
 * Because that's the kind of world I want to live in.  */
void print_rule4(const struct ipt_entry *e,
    struct xtc_handle *h, const char *chain, char **pbuffer)
{
  const struct xt_entry_target *t;
  const char *target_name;

  /* print chain name */
  buffer_printf("-A %s", chain);

  /* Print IP part. */
  print_ip("-s", e->ip.src.s_addr,e->ip.smsk.s_addr,
      e->ip.invflags & IPT_INV_SRCIP, pbuffer);

  print_ip("-d", e->ip.dst.s_addr, e->ip.dmsk.s_addr,
      e->ip.invflags & IPT_INV_DSTIP, pbuffer);

  print_iface('i', e->ip.iniface, e->ip.iniface_mask,
        e->ip.invflags & IPT_INV_VIA_IN, pbuffer);

  print_iface('o', e->ip.outiface, e->ip.outiface_mask,
        e->ip.invflags & IPT_INV_VIA_OUT, pbuffer);

  print_proto(e->ip.proto, e->ip.invflags & XT_INV_PROTO, pbuffer);

  if (e->ip.flags & IPT_F_FRAG)
    buffer_printf("%s -f",
           e->ip.invflags & IPT_INV_FRAG ? " !" : "");

  /* Print matchinfo part */
  if (e->target_offset) {
    IPT_MATCH_ITERATE(e, print_match_save, &e->ip, pbuffer);
  }

  /* Print target name and targinfo part */
  target_name = iptc_get_target(e, h);
  t = ipt_get_target((struct ipt_entry *)e);
  if (t->u.user.name[0]) {
    const struct xtables_target *target =
      xtables_find_target(t->u.user.name, XTF_TRY_LOAD);

    if (!target) {
      fprintf(stderr, "Can't find library for target `%s'\n",
        t->u.user.name);
      return;
    }

    buffer_printf(" -j %s", target_name);
    if (target->save)
      target->save(&e->ip, t);
    else {
      /* If the target size is greater than xt_entry_target
       * there is something to be saved, we just don't know
       * how to print it */
      if (t->u.target_size !=
          sizeof(struct xt_entry_target)) {
        fprintf(stderr, "Target `%s' is missing "
            "save function\n",
          t->u.user.name);
        return;
      }
    }
  } else if (target_name && (*target_name != '\0'))
  buffer_printf(" -%c %s", e->ip.flags & IPT_F_GOTO ? 'g' : 'j', target_name);

  buffer_printf("\n");
}


static int do_output(const char *tablename, char **pbuffer)
{
  struct xtc_handle *h;
  const char *chain = NULL;

  h = iptc_init(tablename);
  if (h == NULL) {
    xtables_load_ko(xtables_modprobe_program, false);
    h = iptc_init(tablename);
  }
  if (!h)
    return -1;

  buffer_printf("*%s\n", tablename);

  /* Dump out chain names first,
   * thereby preventing dependency conflicts */
  for (chain = iptc_first_chain(h);
       chain;
       chain = iptc_next_chain(h)) {

    buffer_printf(":%s ", chain);
    if (iptc_builtin(chain, h)) {
      struct xt_counters count;
      buffer_printf("%s ",
             iptc_get_policy(chain, &count, h));
      buffer_printf("[%llu:%llu]\n", (unsigned long long)count.pcnt, (unsigned long long)count.bcnt);
    } else {
      buffer_printf("- [0:0]\n");
    }
  }

  for (chain = iptc_first_chain(h);
       chain;
       chain = iptc_next_chain(h)) {
    const struct ipt_entry *e;

    /* Dump out rules */
    e = iptc_first_rule(chain, h);
    while(e) {
      print_rule4(e, h, chain, pbuffer);
      e = iptc_next_rule(e, h);
    }
  }

  iptc_free(h);

  return 1;
}
// }}}

static PyObject * iptables_save(PyObject *self, PyObject *args) {
  const char *table_name = NULL;
  char *init_buffer = malloc(1);
  char **pbuffer = &init_buffer;
  init_buffer[0] = 0;

  PyArg_ParseTuple(args, "|s", &table_name);

  if (!table_name) {
    do_output("filter", pbuffer);
    do_output("nat", pbuffer);
  } else {
    do_output(table_name, pbuffer);
  }

  PyObject *result = Py_BuildValue("s", *pbuffer);
  free(*pbuffer);
  return result;
}

static PyMethodDef IptablesMethods[] = {
  {"save",  iptables_save, METH_VARARGS,
    "Similar to iptables-save command."},
  {NULL, NULL, 0, NULL}        /* Sentinel */
};

PyMODINIT_FUNC initiptables(void) {
  xtables_init();
  xtables_set_nfproto(NFPROTO_IPV4);
  (void) Py_InitModule("iptables", IptablesMethods);
}

/*
// setup.py
from distutils.core import setup, Extension

setup(ext_modules=[
    Extension('iptables',
              sources=['iptables_save.c'],
              extra_link_args=['-lip4tc', '-lxtables'])])
// Makefile
.SUFFIXES:

.PHONY: clean all

all:
	python setup.py build_ext --inplace

clean:
	rm -rf build iptables.so
*/