pymin/services/firewall/__init__.py

   1 # vim: set encoding=utf-8 et sw=4 sts=4 :
   2
   3 # TODO See if it's better (more secure) to execute commands via python instead
   4 # of using script templates.
   5
   6 from os import path
   7
   8 from pymin.seqtools import Sequence
   9 from pymin.dispatcher import Handler, handler, HandlerError
  10 from pymin.services.util import Restorable, ConfigWriter, ServiceHandler, \
  11                                 TransactionalHandler
  12
  13 __ALL__ = ('FirewallHandler', 'Error', 'RuleError', 'RuleAlreadyExistsError',
  14            'RuleNotFoundError')
  15
  16 class Error(HandlerError):
  17     r"""
  18     Error(command) -> Error instance :: Base FirewallHandler exception class.
  19
  20     All exceptions raised by the FirewallHandler inherits from this one, so you can
  21     easily catch any FirewallHandler exception.
  22
  23     message - A descriptive error message.
  24     """
  25     pass
  26
  27 class RuleError(Error, KeyError):
  28     r"""
  29     RuleError(rule) -> RuleError instance.
  30
  31     This is the base exception for all rule related errors.
  32     """
  33
  34     def __init__(self, rule):
  35         r"Initialize the object. See class documentation for more info."
  36         self.message = u'Rule error: "%s"' % rule
  37
  38 class RuleAlreadyExistsError(RuleError):
  39     r"""
  40     RuleAlreadyExistsError(rule) -> RuleAlreadyExistsError instance.
  41
  42     This exception is raised when trying to add a rule that already exists.
  43     """
  44
  45     def __init__(self, rule):
  46         r"Initialize the object. See class documentation for more info."
  47         self.message = u'Rule already exists: "%s"' % rule
  48
  49 class RuleNotFoundError(RuleError):
  50     r"""
  51     RuleNotFoundError(rule) -> RuleNotFoundError instance.
  52
  53     This exception is raised when trying to operate on a rule that doesn't
  54     exists.
  55     """
  56
  57     def __init__(self, rule):
  58         r"Initialize the object. See class documentation for more info."
  59         self.message = u'Rule not found: "%s"' % rule
  60
  61
  62 class Rule(Sequence):
  63     r"""Rule(chain, target[, src[, dst[, ...]]]) -> Rule instance.
  64
  65     chain - INPUT, OUTPUT or FORWARD.
  66     target - ACCEPT, REJECT or DROP.
  67     src - Source subnet as IP/mask.
  68     dst - Destination subnet as IP/mask.
  69     protocol - ICMP, UDP, TCP or ALL.
  70     src_port - Source port (only for UDP or TCP protocols).
  71     dst_port - Destination port (only for UDP or TCP protocols).
  72     """
  73
  74     def __init__(self, chain, target, src=None, dst=None, protocol=None,
  75                        src_port=None, dst_port=None):
  76         r"Initialize object, see class documentation for details."
  77         self.chain = chain
  78         self.target = target
  79         self.src = src
  80         self.dst = dst
  81         self.protocol = protocol
  82         # TODO Validate that src_port and dst_port could be not None only
  83         # if the protocol is UDP or TCP
  84         self.src_port = src_port
  85         self.dst_port = dst_port
  86
  87     def update(self, chain=None, target=None, src=None, dst=None, protocol=None,
  88                        src_port=None, dst_port=None):
  89         r"update([chain[, ...]]) -> Update the values of a rule (see Rule doc)."
  90         if chain is not None: self.chain = chain
  91         if target is not None: self.target = target
  92         if src is not None: self.src = src
  93         if dst is not None: self.dst = dst
  94         if protocol is not None: self.protocol = protocol
  95         # TODO Validate that src_port and dst_port could be not None only
  96         # if the protocol is UDP or TCP
  97         if src_port is not None: self.src_port = src_port
  98         if dst_port is not None: self.dst_port = dst_port
  99
 100     def __cmp__(self, other):
 101         r"Compares two Rule objects."
 102         if self.chain == other.chain \
 103                 and self.target == other.target \
 104                 and self.src == other.src \
 105                 and self.dst == other.dst \
 106                 and self.protocol == other.protocol \
 107                 and self.src_port == other.src_port \
 108                 and self.dst_port == other.dst_port:
 109             return 0
 110         return cmp(id(self), id(other))
 111
 112     def as_tuple(self):
 113         r"Return a tuple representing the rule."
 114         return (self.chain, self.target, self.src, self.dst, self.protocol,
 115                     self.src_port, self.dst_port)
 116
 117 class RuleHandler(Handler):
 118     r"""RuleHandler(rules) -> RuleHandler instance :: Handle a list of rules.
 119
 120     This class is a helper for FirewallHandler to do all the work related to rules
 121     administration.
 122
 123     parent - The parent service handler.
 124     """
 125
 126     handler_help = u"Manage firewall rules"
 127
 128     def __init__(self, parent):
 129         r"Initialize the object, see class documentation for details."
 130         self.parent = parent
 131
 132     @handler(u'Add a new rule')
 133     def add(self, *args, **kwargs):
 134         r"add(rule) -> None :: Add a rule to the rules list (see Rule doc)."
 135         rule = Rule(*args, **kwargs)
 136         if rule in self.parent.rules:
 137             raise RuleAlreadyExistsError(rule)
 138         self.parent.rules.append(rule)
 139
 140     @handler(u'Update a rule')
 141     def update(self, index, *args, **kwargs):
 142         r"update(index, rule) -> None :: Update a rule (see Rule doc)."
 143         # TODO check if the modified rule is the same of an existing one
 144         index = int(index) # TODO validation
 145         try:
 146             self.parent.rules[index].update(*args, **kwargs)
 147         except IndexError:
 148             raise RuleNotFoundError(index)
 149
 150     @handler(u'Delete a rule')
 151     def delete(self, index):
 152         r"delete(index) -> Rule :: Delete a rule from the list returning it."
 153         index = int(index) # TODO validation
 154         try:
 155             return self.parent.rules.pop(index)
 156         except IndexError:
 157             raise RuleNotFoundError(index)
 158
 159     @handler(u'Get information about a rule')
 160     def get(self, index):
 161         r"get(rule) -> Rule :: Get all the information about a rule."
 162         index = int(index) # TODO validation
 163         try:
 164             return self.parent.rules[index]
 165         except IndexError:
 166             raise RuleNotFoundError(index)
 167
 168     @handler(u'Get information about all rules')
 169     def show(self):
 170         r"show() -> list of Rules :: List all the complete rules information."
 171         return self.parent.rules
 172
 173
 174 class FirewallHandler(Restorable, ConfigWriter, ServiceHandler,
 175                       TransactionalHandler):
 176     r"""FirewallHandler([pickle_dir[, config_dir]]) -> FirewallHandler instance.
 177
 178     Handles firewall commands using iptables.
 179
 180     pickle_dir - Directory where to write the persistent configuration data.
 181
 182     config_dir - Directory where to store de generated configuration files.
 183
 184     Both defaults to the current working directory.
 185     """
 186
 187     handler_help = u"Manage firewall service"
 188
 189     _persistent_attrs = 'rules'
 190
 191     _restorable_defaults = dict(rules=list())
 192
 193     _config_writer_files = 'iptables.sh'
 194     _config_writer_tpl_dir = path.join(path.dirname(__file__), 'templates')
 195
 196     def __init__(self, pickle_dir='.', config_dir='.'):
 197         r"Initialize the object, see class documentation for details."
 198         self._persistent_dir = pickle_dir
 199         self._config_writer_cfg_dir = config_dir
 200         self._service_start = ('sh', path.join(self._config_writer_cfg_dir,
 201                                         self._config_writer_files))
 202         self._service_stop = ('iptables', '-t', 'filter', '-F')
 203         self._service_restart = self._service_start
 204         self._service_reload = self._service_start
 205         self._config_build_templates()
 206         self._restore()
 207         self.rule = RuleHandler(self)
 208
 209     def _get_config_vars(self, config_file):
 210         return dict(rules=self.rules)
 211
 212
 213 if __name__ == '__main__':
 214
 215     import os
 216
 217     fw_handler = FirewallHandler()
 218
 219     def dump():
 220         print '-' * 80
 221         print 'Rules:'
 222         print fw_handler.rule.show()
 223         print '-' * 80
 224
 225     dump()
 226
 227     fw_handler.rule.add('input','drop','icmp')
 228
 229     fw_handler.rule.update(0, dst='192.168.0.188/32')
 230
 231     fw_handler.rule.add('output','accept', '192.168.1.0/24')
 232
 233     fw_handler.commit()
 234
 235     fw_handler.stop()
 236
 237     dump()
 238
 239     os.system('rm -f *.pkl iptables.sh')
 240