From a7d472af68ebf10283fc4380839e6f30ce8d4a41 Mon Sep 17 00:00:00 2001
From: Thomas Oettli <spacefreak@noop.ch>
Date: Mon, 2 Sep 2019 15:56:12 +0200
Subject: [PATCH] add ability to ignore sender hosts and/or networks

---
 README.md                      |  2 ++
 docs/pyquarantine.conf.example |  6 ++++++
 pyquarantine/__init__.py       | 35 ++++++++++++++++++++++++++++++++--
 3 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 8d9d125..fe7b8a0 100644
--- a/README.md
+++ b/README.md
@@ -49,6 +49,8 @@ The following configuration options are mandatory in each quarantine section:
   SMTP port
 
 The following configuration options are optional in each quarantine section:
+* **ignore_hosts**  
+  Comma-separated list of host and network addresses to be ignored by this quarantine.
 * **reject_reason**  
   Reason to return to the client if action is set to reject.
 
diff --git a/docs/pyquarantine.conf.example b/docs/pyquarantine.conf.example
index 2b03d4c..b13c8bb 100644
--- a/docs/pyquarantine.conf.example
+++ b/docs/pyquarantine.conf.example
@@ -30,6 +30,12 @@ preferred_quarantine_action = last
 
 
 [spam]
+# Option: ignore_hosts
+# Notes:  Set a list of host and network addresses to be ignored by this quarantine.
+#         All the common host/network notations are supported, including IPv6.
+# Value:  [ HOST ]
+#
+ignore_hosts = 10.0.0.0/8,172.16.0.0/12,192.168.0.0/16
 
 # Option: regex
 # Notes:  Set the regular expression to match against email headers.
diff --git a/pyquarantine/__init__.py b/pyquarantine/__init__.py
index 2739838..db8837d 100644
--- a/pyquarantine/__init__.py
+++ b/pyquarantine/__init__.py
@@ -25,7 +25,7 @@ from Milter.utils import parse_addr
 from collections import defaultdict
 from io import BytesIO
 from itertools import groupby
-
+from netaddr import IPAddress, IPNetwork
 from pyquarantine import quarantines
 from pyquarantine import notifications
 from pyquarantine import whitelists
@@ -71,6 +71,20 @@ class QuarantineMilter(Milter.Base):
     def set_configfiles(config_files):
         QuarantineMilter._config_files = config_files
 
+    def connect(self, IPname, family, hostaddr):
+        self.logger.debug("accepted milter connection from {} port {}".format(*hostaddr))
+        ip = IPAddress(hostaddr[0])
+        for quarantine in self.config.copy():
+            for ignore in quarantine["ignore_hosts_list"]:
+                if ip in ignore:
+                    self.logger.debug("host {} is ignored by quarantine {}".format(hostaddr[0], quarantine["name"]))
+                    self.config.remove(quarantine)
+                    break
+            if not self.config:
+                self.logger.debug("host {} is ignored by all quarantines, skip further processing", hostaddr[0])
+                return Milter.ACCEPT
+        return Milter.CONTINUE 
+
     @Milter.noreply
     def envfrom(self, mailfrom, *str):
         self.mailfrom = "@".join(parse_addr(mailfrom)).lower()
@@ -325,7 +339,8 @@ def generate_milter_config(configtest=False, config_files=[]):
 
         # check if optional config options are present in config
         defaults = {
-            "reject_reason": "Message rejected"
+            "reject_reason": "Message rejected",
+            "ignore_hosts": ""
         }
         for option in defaults.keys():
             if option not in config.keys() and \
@@ -391,6 +406,22 @@ def generate_milter_config(configtest=False, config_files=[]):
         else:
             raise RuntimeError("{}: unknown action '{}'".format(quarantine_name, action))
 
+        # create host/network whitelist
+        config["ignore_hosts_list"] = []
+        ignored = set([ p.strip() for p in config["ignore_hosts"].split(",") if p])
+        for ignore in ignored:
+            if not ignore:
+                continue
+            # parse network notation
+            try:
+                net = IPNetwork(ignore)
+            except AddrFormatError as e:
+                raise RuntimeError("error parsing ignore_hosts: {}".format(e))
+            else:
+                config["ignore_hosts_list"].append(net)
+        if config["ignore_hosts_list"]:
+            logger.debug("{}: ignore hosts: {}".format(quarantine_name, ", ".join(ignored)))
+
         milter_config.append(config)
 
     return global_config, milter_config