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 ..admin import HostAdmin
 from firewall.models import Vlan, Domain, Record, Host
 from firewall.models import Vlan, Domain, Record, Host
 from django.forms import ValidationError
 from django.forms import ValidationError
+from ..iptables import IptRule, IptChain, InvalidRuleExcepion
 
 
 
 
 class MockInstance:
 class MockInstance:
@@ -109,3 +110,50 @@ class HostGetHostnameTestCase(TestCase):
         self.r.save()
         self.r.save()
         self.assertEqual(self.h.get_hostname(proto='ipv4', public=True),
         self.assertEqual(self.h.get_hostname(proto='ipv4', public=True),
                          self.r.fqdn)
                          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))