fix file loading and improve quarantine logic

pyquarantine/notify.py

@@ -121,20 +121,21 @@ class EMailNotification(BaseNotification):
         self.from_header = from_header
         self.subject = subject
         try:
-            self.template = open(template, "r").read()
+            with open(template, "r") as fh:
+                self.template = fh.read()
             self.embed_imgs = []
             for img_path in embed_imgs:
-                img = MIMEImage(open(img_path, "rb").read())
+                with open(img_path, "rb") as fh:
+                    img = MIMEImage(fh.read())
                 filename = basename(img_path)
                 img.add_header("Content-ID", f"<{filename}>")
                 self.embed_imgs.append(img)
-
             self.replacement_img = repl_img
             self.strip_images = strip_imgs
 
             if not strip_imgs and repl_img:
-                self.replacement_img = MIMEImage(
-                    open(repl_img, "rb").read())
+                with open(repl_img, "rb") as fh:
+                    self.replacement_img = MIMEImage(fh.read())
                 self.replacement_img.add_header(
                     "Content-ID", "<removed_for_security_reasons>")
 

pyquarantine/storage.py

@@ -177,7 +177,8 @@ class FileMailStorage(BaseMailStorage):
 
         # save mail
         self._save_datafile(datafile, data)
-        logger.info(f"stored message in file {datafile}")
+        logger.debug(f"stored message in file {datafile}")
+        logger.info(f"storaged message with id {storage_id}")
 
         if not self.metadata:
             return storage_id, None, datafile
@@ -393,10 +394,16 @@ class Quarantine:
 
         self._milter_action = None
         if "milter_action" in cfg["args"]:
-            self._milter_action = cfg["args"]["milter_action"]
+            self._milter_action = cfg["args"]["milter_action"].upper()
+            assert self._milter_action in ["ACCEPT", "REJECT"], \
+                f"invalid milter_action '{cfg['args']['milter_action']}'"
+
         self._reason = None
+        if self._milter_action == "REJECT":
             if "reject_reason" in cfg["args"]:
                 self._reason = cfg["args"]["reject_reason"]
+            else:
+                self._reason = "Message rejected"
 
     def __str__(self):
         cfg = []
@@ -483,19 +490,17 @@ class Quarantine:
     def execute(self, milter):
         logger = CustomLogger(
             self.logger, {"name": self.cfg["name"], "qid": milter.qid})
+        rcpts = milter.msginfo["rcpts"]
         wl_rcpts = []
         if self._whitelist:
             wl_rcpts = self._whitelist.get_wl_rcpts(
-                milter.msginfo["mailfrom"], milter.msginfo["rcpts"], logger)
+                milter.msginfo["mailfrom"], rcpts, logger)
+            if wl_rcpts:
                 logger.info(f"whitelisted recipients: {wl_rcpts}")
-
-        rcpts = [
-            rcpt for rcpt in milter.msginfo["rcpts"] if rcpt not in wl_rcpts]
-
+            rcpts = [rcpt for rcpt in rcpts if rcpt not in wl_rcpts]
             if not rcpts:
                 # all recipients whitelisted
                 return
-
             logger.info(f"add to quarantine for recipients: {rcpts}")
             milter.msginfo["rcpts"] = rcpts
 
@@ -505,7 +510,8 @@ class Quarantine:
             self._notification.execute(milter)
 
         milter.msginfo["rcpts"].extend(wl_rcpts)
-        milter.delrcpt(rcpts)
 
-        if self._milter_action is not None and not milter.msginfo["rcpts"]:
+        if self._milter_action is not None:
+            milter.delrcpt(rcpts)
+            if not milter.msginfo["rcpts"]:
                 return (self._milter_action, self._reason)