1 files changed, 200 insertions, 0 deletions
diff --git a/lib/ipset.py b/lib/ipset.py
new file mode 100644
index 0000000..2ff4fbb
--- /dev/null
+++ b/lib/ipset.py
@@ -0,0 +1,200 @@
+#!/usr/bin/python
+#
+# Copyright 2013 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Ipset iptables generator.  This is a subclass of Iptables generator.
+
+ipset is a system inside the Linux kernel, which can very efficiently store
+and match IPv4 and IPv6 addresses. This can be used to dramatically increase
+performace of iptables firewall.
+
+"""
+
+__author__ = 'vklimovs@google.com (Vjaceslavs Klimovs)'
+
+from string import Template
+
+import iptables
+import nacaddr
+
+
+class Error(Exception):
+  pass
+
+
+class Term(iptables.Term):
+  """Single Ipset term representation."""
+
+  _PLATFORM = 'ipset'
+  _SET_MAX_LENGTH = 31
+  _POSTJUMP_FORMAT = None
+  _PREJUMP_FORMAT = None
+  _TERM_FORMAT = None
+  _COMMENT_FORMAT = Template('-A $filter -m comment --comment "$comment"')
+  _FILTER_TOP_FORMAT = Template('-A $filter')
+
+  def __init__(self, *args, **kwargs):
+    super(Term, self).__init__(*args, **kwargs)
+    # This stores tuples of set name and set contents, keyed by direction.
+    # For example:
+    # { 'src': ('term_name', [ipaddr object, ipaddr object]),
+    #  'dst': ('term_name', [ipaddr object, ipaddr object]) }
+    self.addr_sets = dict()
+
+  def _CalculateAddresses(self, src_addr_list, src_ex_addr_list,
+                          dst_addr_list, dst_ex_addr_list):
+    """Calculate source and destination address list for a term.
+
+    Since ipset is very efficient at matching large number of
+    addresses, we never return eny exclude addresses. Instead
+    least positive match is calculated for both source and destination
+    addresses.
+
+    For source and destination address list, three cases are possible.
+    First case is when there is no addresses. In that case we return
+    _all_ips.
+    Second case is when there is strictly one address. In that case,
+    we optimize by not generating a set, and it's then the only
+    element of returned set.
+    Third case case is when there is more than one address in a set.
+    In that case we generate a set and also return _all_ips. Note the
+    difference to the first case where no set is actually generated.
+
+    Args:
+      src_addr_list: source address list of the term.
+      src_ex_addr_list: source address exclude list of the term.
+      dst_addr_list: destination address list of the term.
+      dst_ex_addr_list: destination address exclude list of the term.
+
+    Returns:
+      tuple containing source address list, source exclude address list,
+      destination address list, destination exclude address list in
+      that order.
+
+    """
+    if not src_addr_list:
+      src_addr_list = [self._all_ips]
+    src_addr_list = [src_addr for src_addr in src_addr_list if
+                     src_addr.version == self.AF_MAP[self.af]]
+    if src_ex_addr_list:
+      src_ex_addr_list = [src_ex_addr for src_ex_addr in src_ex_addr_list if
+                          src_ex_addr.version == self.AF_MAP[self.af]]
+      src_addr_list = nacaddr.ExcludeAddrs(src_addr_list, src_ex_addr_list)
+    if len(src_addr_list) > 1:
+      set_name = self._GenerateSetName(self.term.name, 'src')
+      self.addr_sets['src'] = (set_name, src_addr_list)
+      src_addr_list = [self._all_ips]
+
+    if not dst_addr_list:
+      dst_addr_list = [self._all_ips]
+    dst_addr_list = [dst_addr for dst_addr in dst_addr_list if
+                     dst_addr.version == self.AF_MAP[self.af]]
+    if dst_ex_addr_list:
+      dst_ex_addr_list = [dst_ex_addr for dst_ex_addr in dst_ex_addr_list if
+                          dst_ex_addr.version == self.AF_MAP[self.af]]
+      dst_addr_list = nacaddr.ExcludeAddrs(dst_addr_list, dst_ex_addr_list)
+    if len(dst_addr_list) > 1:
+      set_name = self._GenerateSetName(self.term.name, 'dst')
+      self.addr_sets['dst'] = (set_name, dst_addr_list)
+      dst_addr_list = [self._all_ips]
+    return (src_addr_list, [], dst_addr_list, [])
+
+  def _GenerateAddressStatement(self, src_addr, dst_addr):
+    """Return the address section of an individual iptables rule.
+
+    See _CalculateAddresses documentation. Three cases are possible here,
+    and they map directly to cases in _CalculateAddresses.
+    First, there can be no addresses for a direction (value is _all_ips then)
+    In that case we return empty string.
+    Second there can be stricly one address. In that case we return single
+    address match (-s or -d).
+    Third case, is when the value is _all_ips but also the set for particular
+    direction is present. That's when we return a set match.
+
+    Args:
+      src_addr: source address of the rule.
+      dst_addr: destination address of the rule.
+
+    Returns:
+      tuple containing source and destination address statement, in
+      that order.
+
+    """
+    src_addr_stmt = ''
+    dst_addr_stmt = ''
+    if src_addr and dst_addr:
+      if src_addr == self._all_ips:
+        if 'src' in self.addr_sets:
+          src_addr_stmt = ('-m set --set %s src' % self.addr_sets['src'][0])
+      else:
+        src_addr_stmt = '-s %s/%d' % (src_addr.ip, src_addr.prefixlen)
+      if dst_addr == self._all_ips:
+        if 'dst' in self.addr_sets:
+          dst_addr_stmt = ('-m set --set %s dst' % self.addr_sets['dst'][0])
+      else:
+        dst_addr_stmt = '-d %s/%d' % (dst_addr.ip, dst_addr.prefixlen)
+    return (src_addr_stmt, dst_addr_stmt)
+
+  def _GenerateSetName(self, term_name, suffix):
+    if self.af == 'inet6':
+      suffix += '-v6'
+    if len(term_name) + len(suffix) + 1 > self._SET_MAX_LENGTH:
+      term_name = term_name[:self._SET_MAX_LENGTH -
+                            (len(term_name) + len(suffix) + 1)]
+    return term_name + '-' + suffix
+
+
+class Ipset(iptables.Iptables):
+  """Ipset generator."""
+  _PLATFORM = 'ipset'
+  _SET_TYPE = 'hash:net'
+  _SUFFIX = '.ips'
+  _TERM = Term
+
+  def __str__(self):
+    # Actual rendering happens in __str__, so it has to be called
+    # before we do set specific part.
+    iptables_output = iptables.Iptables.__str__(self)
+    output = []
+    for (_, _, _, _, terms) in self.iptables_policies:
+      for term in terms:
+        output.extend(self._GenerateSetConfig(term))
+    output.append(iptables_output)
+    return '\n'.join(output)
+
+  def _GenerateSetConfig(self, term):
+    """Generate set configuration for supplied term.
+
+    Args:
+      term: input term.
+
+    Returns:
+      string that is configuration of supplied term.
+
+    """
+    output = []
+    for direction in sorted(term.addr_sets, reverse=True):
+      set_hashsize = 2 ** len(term.addr_sets[direction][1]).bit_length()
+      set_maxelem = 2 ** len(term.addr_sets[direction][1]).bit_length()
+      output.append('create %s %s family %s hashsize %i maxelem %i' %
+                    (term.addr_sets[direction][0],
+                     self._SET_TYPE,
+                     term.af,
+                     set_hashsize,
+                     set_maxelem))
+      for address in term.addr_sets[direction][1]:
+        output.append('add %s %s' % (term.addr_sets[direction][0], address))
+    return output