fix bug in quarantine whitelist

README.md

@@ -120,7 +120,7 @@ Config options for conditions objects:
 * **headers** (optional)  
   Matches e-mails for which all regular expressions in the given list are matching at least one e-mail header.
 * **whitelist** (optional)  
-  Matches e-mails for which the whitelist has no entry for the envelope-from and envelope-to address combination, see section [Whitelist](#Whitelist).
+  Matches e-mails for which the whitelist has no entry for the envelope-from and envelope-to address combination, see section [Whitelist](#Whitelist) for details.
 * **var** (optional)  
   Matches e-mails for which a previous action or condition has set the given metavariable.
 * **metavar** (optional)  
@@ -210,14 +210,14 @@ Options:
 Quarantine e-mail.  
 Options:
 * **store**  
-  Options for e-mail storage, see action **store** in section [Action types](#Action-types).  
+  Options for e-mail storage, see action [store](#store).  
   If the option **metadata** is not specificall set for this storage, it will be set to true.
 * **smtp_host**  
   SMTP host used to release e-mails from quarantine.
 * **smtp_port**  
   SMTP port used to release e-mails from quarantine.
 * **notify** (optional)  
-  Options for e-mail notifications, see action **notify** in section [Action types](#Action-types).
+  Options for e-mail notifications, see action [notify](#notify).
 * **milter_action** (optional)  
   Milter action to perform. If set, no further rules or actions will be processed.  
   Please think carefully what you set here or your MTA may do something you do not want it to do.  
@@ -231,7 +231,7 @@ Options:
 * **reject_reason** (optional, default: "Message rejected")  
   Reject message sent to MTA if milter_action is set to reject.
 * **whitelist** (optional)  
-  Options for a whitelist, see **whitelist** in section [Conditions](#Conditions).
+  Options for a whitelist object, see section [Whitelist](#Whitelist).
 
 ### Storage types
 Available storage types:

distribution/pypi/build.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 set -e
-
+set -x
 PYTHON=$(which python)
 
 script_dir=$(dirname "$(readlink -f -- "$BASH_SOURCE")")

pyquarantine/__init__.py

@@ -27,7 +27,7 @@ __all__ = [
     "whitelist",
     "QuarantineMilter"]
 
-__version__ = "2.0.0"
+__version__ = "2.0.2"
 
 from pyquarantine import _runtime_patches
 

pyquarantine/base.py

@@ -65,6 +65,133 @@ class MilterMessage(MIMEPart):
 
         self._headers = newheaders
 
+    def _find_body_parent(self, part, preferencelist, parent=None):
+        if part.is_attachment():
+            return
+        maintype, subtype = part.get_content_type().split("/")
+        if maintype == "text":
+            if subtype in preferencelist:
+                yield(preferencelist.index(subtype), parent)
+            return
+        if maintype != "multipart" or not self.is_multipart():
+            return
+        if subtype != "related":
+            for subpart in part.iter_parts():
+                yield from self._find_body_parent(
+                    subpart, preferencelist, part)
+            return
+        if 'related' in preferencelist:
+            yield(preferencelist.index('related'), parent)
+        candidate = None
+        start = part.get_param('start')
+        if start:
+            for subpart in part.iter_parts():
+                if subpart['content-id'] == start:
+                    candidate = subpart
+                    break
+        if candidate is None:
+            subparts = part.get_payload()
+            candidate = subparts[0] if subparts else None
+        if candidate is not None:
+            yield from self._find_body_parent(candidate, preferencelist, part)
+
+    def get_body_parent(self, preferencelist=("related", "html", "plain")):
+        best_prio = len(preferencelist)
+        body_parent = None
+        for prio, parent in self._find_body_parent(self, preferencelist):
+            if prio < best_prio:
+                best_prio = prio
+                body_parent = parent
+                if prio == 0:
+                    break
+        return body_parent
+
+    def get_body_content(self, pref):
+        part = None
+        content = None
+        if not self.is_multipart() and \
+                self.get_content_type() == f"text/{pref}":
+            part = self
+        else:
+            part = self.get_body(preferencelist=(pref))
+
+        if part is not None:
+            content = part.get_content()
+
+        return (part, content)
+
+    def set_body(self, text_content=None, html_content=None):
+        parent = self.get_body_parent() or self
+        if "Content-Type" not in parent:
+            # set Content-Type header if not present, otherwise
+            # make_alternative and make_mixed skip the payload
+            parent["Content-Type"] = parent.get_content_type()
+
+        maintype, subtype = parent.get_content_type().split("/")
+        if not parent.is_multipart() or maintype != "multipart":
+            if maintype == "text" and subtype in ("html", "plain"):
+                parent.make_alternative()
+                maintype, subtype = ("multipart", "alternative")
+            else:
+                parent.make_mixed()
+                maintype, subtype = ("multipart", "mixed")
+
+        text_body = parent.get_body(preferencelist=("plain"))
+        html_body = parent.get_body(preferencelist=("html"))
+
+        if text_content is not None:
+            if text_body:
+                text_body.set_content(text_content)
+            else:
+                if not html_body or subtype == "alternative":
+                    inject_body_part(parent, text_content)
+                else:
+                    html_body.add_alternative(text_content)
+                text_body = parent.get_body(preferencelist=("plain"))
+
+        if html_content is not None:
+            if html_body:
+                html_body.set_content(html_content, subtype="html")
+            else:
+                if not text_body or subtype == "alternative":
+                    inject_body_part(parent, html_content, subtype="html")
+                else:
+                    text_body.add_alternative(html_content, subtype="html")
+
+
+def inject_body_part(part, content, subtype="plain"):
+    parts = []
+    text_body = None
+    if subtype == "html":
+        text_body = part.get_body(preferencelist=("plain"))
+
+    for p in part.iter_parts():
+        if text_body and p == text_body:
+            continue
+        parts.append(p)
+
+    boundary = part.get_boundary()
+    p_subtype = part.get_content_subtype()
+    part.clear_content()
+    if text_body:
+        part.set_content(content)
+        part.add_alternative(content, subtype=subtype)
+    else:
+        part.set_content(content, subtype=subtype)
+
+    if part.get_content_subtype() != p_subtype:
+        if p_subtype == "alternative":
+            part.make_alternative()
+        elif p_subtype == "related":
+            part.make_related()
+        else:
+            part.make_mixed()
+
+    if boundary:
+        part.set_boundary(boundary)
+    for p in parts:
+        part.attach(p)
+
 
 def replace_illegal_chars(string):
     """Remove illegal characters from header values."""

pyquarantine/modify.py

@@ -183,26 +183,6 @@ def _has_content_before_body_tag(soup):
     return False
 
 
-def _inject_body(milter, text_content="", html_content=None):
-    if not milter.msg.is_multipart():
-        milter.msg.make_mixed()
-
-    attachments = []
-    for attachment in milter.msg.iter_attachments():
-        if "content-disposition" not in attachment:
-            attachment["Content-Disposition"] = "attachment"
-        attachments.append(attachment)
-
-    milter.msg.clear_content()
-    milter.msg.set_content(text_content)
-    if html_content is not None:
-        milter.msg.add_alternative(html_content, subtype="html")
-    milter.msg.make_mixed()
-
-    for attachment in attachments:
-        milter.msg.attach(attachment)
-
-
 def _wrap_message(milter):
     attachment = MIMEPart(policy=SMTP)
     attachment.set_content(milter.msg_as_bytes(),
@@ -248,29 +228,28 @@ class AddDisclaimer:
         self.pretend = pretend
 
     def patch_message_body(self, milter, logger):
-        text_body, text_content = _get_body_content(milter.msg, "plain")
-        html_body, html_content = _get_body_content(milter.msg, "html")
+        text_body, text_content = milter.msg.get_body_content("plain")
+        html_body, html_content = milter.msg.get_body_content("html")
 
         if text_content is None and html_content is None:
-            logger.info("message contains only attachment(s), "
-                        "inject empty body")
+            logger.info("message contains no body, inject it")
             if self.add_html_body:
-                _inject_body(milter, "", "")
-                html_body, html_content = _get_body_content(milter.msg, "html")
+                milter.msg.set_body("", "")
+                html_body, html_content = milter.msg.get_body_content("html")
             else:
-                _inject_body(milter, "")
-            text_body, text_content = _get_body_content(milter.msg, "plain")
+                milter.msg.set_body("")
+            text_body, text_content = milter.msg.get_body_content("plain")
 
         if html_content is None and self.add_html_body:
-            logger.info("inject html body part generated from plain body")
+            logger.info("inject html body based on plain body")
             header = '<meta http-equiv="Content-Type" content="text/html; ' \
                      'charset=utf-8">'
             html_text = re.sub(r"^(.*)$", r"\1<br/>",
                                escape(text_content, quote=False),
                                flags=re.MULTILINE)
-            _inject_body(milter, text_content, f"{header}{html_text}")
-            text_body, text_content = _get_body_content(milter.msg, "plain")
-            html_body, html_content = _get_body_content(milter.msg, "html")
+            milter.msg.set_body(None, f"{header}{html_text}")
+            text_body, text_content = milter.msg.get_body_content("plain")
+            html_body, html_content = milter.msg.get_body_content("html")
 
         variables = defaultdict(str, milter.msginfo["vars"])
         variables["ENVELOPE_FROM"] = escape(

pyquarantine/run.py

@@ -149,7 +149,7 @@ def main():
 
     # increase the recursion limit so that BeautifulSoup can
     # parse larger html content
-    sys.setrecursionlimit(2000)
+    sys.setrecursionlimit(4000)
 
     rc = 0
     try:

pyquarantine/storage.py

@@ -549,7 +549,7 @@ class Quarantine:
             if not rcpts:
                 # all recipients whitelisted
                 return
-            milter.msginfo["rcpts"] = rcpts
+            milter.msginfo["rcpts"] = rcpts.copy()
 
         if self._milter_action in ["REJECT", "DISCARD"]:
             logger.info(f"quarantine message for {rcpts}")
@@ -563,6 +563,5 @@ class Quarantine:
 
         if self._milter_action is not None:
             milter.delrcpt(rcpts)
-            if self._milter_action in ["ACCEPT", "REJECT"] and \
-                    not milter.msginfo["rcpts"]:
+            if not milter.msginfo["rcpts"]:
                 return (self._milter_action, self._reason)