fw: refactor fw.py

circle/firewall/iptables.py

+import logging
+import re
+from collections import OrderedDict
+
+logger = logging.getLogger()
+
+ipv4_re = re.compile(
+    r'^(25[0-5]|2[0-4]\d|[0-1]?\d?\d)(\.(25[0-5]|2[0-4]\d|[0-1]?\d?\d)){3}')
+
+
+class InvalidRuleExcepion(Exception):
+    pass
+
+
+class IptRule(object):
+
+    def __init__(self, priority=1000, action=None, src=None, dst=None,
+                 proto=None, sport=None, dport=None, extra=None):
+        if proto not in ['tcp', 'udp', 'icmp', None]:
+            raise InvalidRuleExcepion()
+        if proto not in ['tcp', 'udp'] and (sport is not None or
+                                            dport is not None):
+            raise InvalidRuleExcepion()
+
+        self.priority = int(priority)
+        self.action = action
+
+        (self.src4, self.src6) = (None, None)
+        if isinstance(src, tuple):
+            (self.src4, self.src6) = src
+        (self.dst4, self.dst6) = (None, None)
+        if isinstance(dst, tuple):
+            (self.dst4, self.dst6) = dst
+
+        self.proto = proto
+        self.sport = sport
+        self.dport = dport
+
+        self.extra = extra
+        self.ipv4_only = extra and bool(ipv4_re.search(extra))
+
+    def __hash__(self):
+        return hash(frozenset(self.__dict__.items()))
+
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+    def __lt__(self, other):
+        return self.priority < other.priority
+
+    def __repr__(self):
+        return '<IptRule: @%d %s >' % (self.priority, self.compile())
+
+    def __unicode__(self):
+        return self.__repr__()
+
+    def compile(self, proto='ipv4'):
+        opts = OrderedDict([('src4' if proto == 'ipv4' else 'src6', '-s %s'),
+                            ('dst4' if proto == 'ipv4' else 'dst6', '-d %s'),
+                            ('proto', '-p %s'),
+                            ('sport', '--sport %s'),
+                            ('dport', '--dport %s'),
+                            ('extra', '%s'),
+                            ('action', '-g %s')])
+        params = [opts[param] % getattr(self, param)
+                  for param in opts
+                  if getattr(self, param) is not None]
+        return ' '.join(params)
+
+
+class IptChain(object):
+    builtin_chains = ('FORWARD', 'INPUT', 'OUTPUT', 'PREROUTING',
+                      'POSTROUTING')
+
+    def __init__(self, name):
+        self.rules = set()
+        self.name = name
+
+    def add(self, *args, **kwargs):
+        for rule in args:
+            self.rules.add(rule)
+
+    def sort(self):
+        return sorted(list(self.rules))
+
+    def __len__(self):
+        return len(self.rules)
+
+    def __repr__(self):
+        return '<IptChain: %s %s>' % (self.name, self.rules)
+
+    def __unicode__(self):
+        return self.__repr__()
+
+    def compile(self, proto='ipv4'):
+        assert proto in ('ipv4', 'ipv6')
+        prefix = '-A %s ' % self.name
+        return '\n'.join([prefix + rule.compile(proto)
+                          for rule in self.sort()
+                          if not (proto == 'ipv6' and rule.ipv4_only)])

circle/firewall/templates/firewall/iptables.conf

+{% if nat %}
+*nat
+:PREROUTING ACCEPT [0:0]
+:INPUT ACCEPT [0:0]
+:OUTPUT ACCEPT [0:0]
+:POSTROUTING ACCEPT [0:0]
+{% for chain in nat %}
+{{ chain.compile }}
+{% endfor %}
+COMMIT
+{% endif %}
+
+*filter
+:INPUT DROP [0:0]
+:FORWARD DROP [0:0]
+:OUTPUT DROP [0:0]
+
+# initialize logging
+-N LOG_DROP
+# windows port scan are silently dropped
+-A LOG_DROP -p tcp --dport 445 -j DROP
+-A LOG_DROP -p udp --dport 137 -j DROP
+-A LOG_DROP -j LOG --log-level 7 --log-prefix "[ipt][drop]"
+-A LOG_DROP -j DROP
+-N LOG_ACC
+-A LOG_ACC -j LOG --log-level 7 --log-prefix "[ipt][isok]"
+-A LOG_ACC -j ACCEPT
+
+# initialize FORWARD chain
+-A FORWARD -m set --match-set blacklist src,dst -j DROP
+-A FORWARD -m state --state INVALID -g LOG_DROP
+-A FORWARD -m state --state ESTABLISHED,RELATED -j ACCEPT
+-A FORWARD -p icmp --icmp-type echo-request -g LOG_ACC
+
+# initialize INPUT chain
+-A INPUT -m set --match-set blacklist src -j DROP
+-A INPUT -m state --state INVALID -g LOG_DROP
+-A INPUT -i lo -j ACCEPT
+-A INPUT -m state --state ESTABLISHED,RELATED -j ACCEPT
+
+# initialize OUTPUT chain
+-A OUTPUT -m state --state INVALID -g LOG_DROP
+-A OUTPUT -o lo -j ACCEPT
+-A OUTPUT -m state --state ESTABLISHED,RELATED -j ACCEPT
+
+{% for chain in filter %}
+{% if chain.name not in chain.builtin_chains %}-N {{ chain.name }}{% endif %}
+{{ chain.compile }}
+{% endfor %}
+
+# close all chains
+-A FORWARD -g LOG_DROP
+-A INPUT -g LOG_DROP
+-A OUTPUT -g LOG_DROP
+COMMIT

circle/firewall/tests/test_firewall.py

@@ -5,6 +5,7 @@ from django.contrib.auth.models import User
 from ..admin import HostAdmin
 from firewall.models import Vlan, Domain, Record, Host
 from django.forms import ValidationError
+from ..iptables import IptRule, IptChain, InvalidRuleExcepion
 
 
 class MockInstance:
@@ -109,3 +110,50 @@ class HostGetHostnameTestCase(TestCase):
         self.r.save()
         self.assertEqual(self.h.get_hostname(proto='ipv4', public=True),
                          self.r.fqdn)
+
+
+class IptablesTestCase(TestCase):
+    def setUp(self):
+        self.r = [IptRule(priority=4, action='ACCEPT',
+                          src=('127.0.0.4', None)),
+                  IptRule(priority=4, action='ACCEPT',
+                          src=('127.0.0.4', None)),
+                  IptRule(priority=2, action='ACCEPT',
+                          dst=('127.0.0.2', None),
+                          extra='-p icmp'),
+                  IptRule(priority=6, action='ACCEPT',
+                          dst=('127.0.0.6', None),
+                          proto='tcp', dport=80),
+                  IptRule(priority=1, action='ACCEPT',
+                          dst=('127.0.0.1', None),
+                          proto='udp', dport=53),
+                  IptRule(priority=5, action='ACCEPT',
+                          dst=('127.0.0.5', None),
+                          proto='tcp', dport=443),
+                  IptRule(priority=2, action='ACCEPT',
+                          dst=('127.0.0.2', None),
+                          proto='icmp'),
+                  IptRule(priority=6, action='ACCEPT',
+                          dst=('127.0.0.6', None),
+                          proto='tcp', dport='1337')]
+
+    def test_chain_add(self):
+        ch = IptChain(name='test')
+        ch.add(*self.r)
+        self.assertEqual(len(ch), len(self.r) - 1)
+
+    def test_rule_compile_ok(self):
+        self.assertEqual(self.r[5].compile(),
+                         '-d 127.0.0.5 -p tcp --dport 443 -g ACCEPT')
+
+    def test_rule_compile_fail(self):
+        self.assertRaises(InvalidRuleExcepion,
+                          IptRule, **{'priority': 5, 'action': 'ACCEPT',
+                                      'dst': '127.0.0.5',
+                                      'proto': 'icmp', 'dport': 443})
+
+    def test_chain_compile(self):
+        ch = IptChain(name='test')
+        ch.add(*self.r)
+        compiled = ch.compile()
+        self.assertEqual(len(compiled.splitlines()), len(ch))