Cleanup source code

pyheadermilter/__init__.py

@@ -29,38 +29,48 @@ from netaddr import IPAddress, IPNetwork, AddrFormatError
 
 
 class HeaderRule:
-    def __init__(self, name, action, header, search=None, value=None, ignore_hosts=[], only_hosts=[], log=True):
+    """HeaderRule to implement a rule to apply on e-mail headers."""
+
+    def __init__(self, name, action, header, search="", value="", ignore_hosts=[], only_hosts=[], log=True):
         self.logger = logging.getLogger(__name__)
         self.name = name
         self._action = action
-        if action == "add":
+        self.header = header
-            self.header = header
+        self.search = search
-        else:
+        self.value = value
+        self.ignore_hosts = ignore_hosts
+        self.only_hosts = only_hosts
+        self.log = log
+
+        if action in ["del", "mod"]:
+            # compile header regex
             try:
                 self.header = re.compile(header, re.MULTILINE + re.DOTALL + re.IGNORECASE)
             except re.error as e:
                 raise RuntimeError("unable to parse option 'header' of rule '{}': {}".format(name, e))
-        if action == "mod":
+
-            try:
+            if action == "mod":
-                self.search = re.compile(search, re.MULTILINE + re.DOTALL + re.IGNORECASE)
+                # compile search regex
-            except re.error as e:
+                try:
-                raise RuntimeError("unable to parse option 'search' of rule '{}': {}".format(name, e))
+                    self.search = re.compile(search, re.MULTILINE + re.DOTALL + re.IGNORECASE)
-        else:
+                except re.error as e:
-            self.search = search
+                    raise RuntimeError("unable to parse option 'search' of rule '{}': {}".format(name, e))
-        self.value = value
+
-        self.ignore_hosts = []
+        if action in ["add", "mod"] and not value:
+            raise RuntimeError("value of option 'value' is empty")
+
+        # replace strings in ignore_hosts and only_hosts with IPNetwork instances
         try:
-            for ignore in ignore_hosts:
+            for index, ignore in enumerate(ignore_hosts):
-                self.ignore_hosts.append(IPNetwork(ignore))
+                self.ignore_hosts[index] = IPNetwork(ignore)
         except AddrFormatError as e:
             raise RuntimeError("unable to parse option 'ignore_hosts' of rule '{}': {}".format(name, e))
-        self.only_hosts = []
+
         try:
-            for only in only_hosts:
+            for index, only in enumerate(only_hosts):
-                self.only_hosts.append(IPNetwork(only))
+                self.only_hosts[index] = IPNetwork(only)
         except AddrFormatError as e:
             raise RuntimeError("unable to parse option 'only_hosts' of rule '{}': {}".format(name, e))
-        self.log = log
 
     def get_action(self):
         return self._action
@@ -71,39 +81,51 @@ class HeaderRule:
     def ignore_host(self, host):
         ip = IPAddress(host)
         ignore = False
+
+        # check if host matches ignore_hosts
         for ignored in self.ignore_hosts:
             if ip in ignored:
                 ignore = True
                 break
+
         if not ignore and self.only_hosts:
+            # host does not match ignore_hosts, check if it matches only_hosts
             ignore = True
             for only in self.only_hosts:
                 if ip in only:
                     ignore = False
                     break
+
         if ignore:
             self.logger.debug("host {} is ignored by rule {}".format(host, self.name))
         return ignore
 
     def execute(self, headers):
-        modified = []
+        """Execute rule on given headers and return list with modified headers."""
         if self._action == "add":
-            modified.append((self.header, self.value, 0, 1))
+            return [(self.header, self.value, 0, 1)]
-        else:
+
-            index = 0
+        modified = []
-            occurrences = {}
+        index = 0
-            for name, value in headers:
+        occurrences = {}
-                if name not in occurrences.keys():
+
-                    occurrences[name] = 1
+        # iterate headers
+        for name, value in headers:
+            # keep track of the occurrence of each header, needed by Milter.Base.chgheader
+            if name not in occurrences.keys():
+                occurrences[name] = 1
+            else:
+                occurrences[name] += 1
+
+            # check if header line matches regex
+            if self.header.search("{}: {}".format(name, value)):
+                if self._action == "del":
+                    # set an empty value to delete the header
+                    value = ""
                 else:
-                    occurrences[name] += 1
+                    value = self.search.sub(self.value, value)
-                if self.header.search("{}: {}".format(name, value)):
+                modified.append((name, value, index, occurrences[name]))
-                    if self._action == "del":
+            index += 1
-                        value = ""
-                    else:
-                        value = self.search.sub(self.value, value)
-                    modified.append((name, value, index, occurrences[name]))
-                index += 1
         return modified
 
 
@@ -124,9 +146,12 @@ class HeaderMilter(Milter.Base):
     def connect(self, IPname, family, hostaddr):
         self.logger.debug("accepted milter connection from {} port {}".format(*hostaddr))
         ip = IPAddress(hostaddr[0])
+
+        # remove rules which ignore this host
         for rule in self.rules.copy():
             if rule.ignore_host(ip):
                 self.rules.remove(rule)
+
         if not self.rules:
             self.logger.debug("host {} is ignored by all rules, skip further processing".format(hostaddr[0]))
             return Milter.ACCEPT
@@ -147,20 +172,21 @@ class HeaderMilter(Milter.Base):
     def eom(self):
         try:
             for rule in self.rules:
-                action = rule.get_action()
-                log = rule.log_modification()
                 self.logger.debug("{}: executing rule '{}'".format(self.queueid, rule.name))
                 modified = rule.execute(self.headers)
+                action = rule.get_action()
+                log = rule.log_modification()
+
                 for name, value, index, occurrence in modified:
                     if action == "add":
                         if log:
                             self.logger.info("{}: add: header: {}".format(self.queueid, "{}: {}".format(name, value)[0:50]))
                         else:
                             self.logger.debug("{}: add: header: {}".format(self.queueid, "{}: {}".format(name, value)[0:50]))
+
                         self.headers.insert(0, (name, value))
                         self.addheader(name, value, 1)
                     else:
-                        self.chgheader(name, occurrence, value)
                         if action == "mod":
                             if log:
                                 self.logger.info("{}: modify: header: {}".format(self.queueid, "{}: {}".format(name, self.headers[index][1])[0:50]))
@@ -173,6 +199,7 @@ class HeaderMilter(Milter.Base):
                             else:
                                 self.logger.debug("{}: delete: header (occ. {}): {}".format(self.queueid, occurrence, "{}: {}".format(name, self.headers[index][1])[0:50]))
                             del self.headers[index]
+                        self.chgheader(name, occurrence, value)
             return Milter.ACCEPT
         except Exception as e:
             self.logger.exception("an exception occured in eom function: {}".format(e))
@@ -293,6 +320,7 @@ def main():
                 config["log"] = False
             else:
                 raise RuntimeError("invalid value specified for option 'log' for rule '{}'".format(rule_name))
+
             # add rule
             logging.debug("adding rule '{}'".format(rule_name))
             rules.append(HeaderRule(name=rule_name, **config))