improve error handling

2021-09-10 15:34:58 +02:00
parent e632f0d511
commit 2b73e43d2f
4 changed files with 64 additions and 85 deletions
--- a/pymodmilter/action.py
+++ b/pymodmilter/action.py
@@ -17,9 +17,6 @@ __all__ = [
    "Action"]

 import os
-import re
-
-from bs4 import BeautifulSoup

 from pymodmilter import CustomLogger, BaseConfig
 from pymodmilter.conditions import ConditionsConfig, Conditions
@@ -60,70 +57,33 @@ class ActionConfig(BaseConfig):
        if self["type"] == "add_header":
            self["class"] = modify.AddHeader
            self["headersonly"] = True
-
-            if "field" not in cfg and "header" in cfg:
-                cfg["field"] = cfg["header"]
-
-            self.add_string_arg(cfg, ("field", "value"))
-
+            self.add_string_arg(cfg, ["field", "value"])
        elif self["type"] == "mod_header":
            self["class"] = modify.ModHeader
            self["headersonly"] = True
-
-            if "field" not in cfg and "header" in cfg:
-                cfg["field"] = cfg["header"]
-
            args = ["field", "value"]
            if "search" in cfg:
                args.append("search")

-            for arg in args:
-                self.add_string_arg(cfg, arg)
-                if arg in ("field", "search"):
-                    try:
-                        self["args"][arg] = re.compile(
-                            self["args"][arg],
-                            re.MULTILINE + re.DOTALL + re.IGNORECASE)
-                    except re.error as e:
-                        raise ValueError(f"{self['name']}: {arg}: {e}")
-
+            self.add_string_arg(cfg, args)
        elif self["type"] == "del_header":
            self["class"] = modify.DelHeader
            self["headersonly"] = True
-
-            if "field" not in cfg and "header" in cfg:
-                cfg["field"] = cfg["header"]
-
            args = ["field"]
            if "value" in cfg:
                args.append("value")

-            for arg in args:
-                self.add_string_arg(cfg, arg)
-                try:
-                    self["args"][arg] = re.compile(
-                        self["args"][arg],
-                        re.MULTILINE + re.DOTALL + re.IGNORECASE)
-                except re.error as e:
-                    raise ValueError(f"{self['name']}: {arg}: {e}")
-
+            self.add_string_arg(cfg, args)
        elif self["type"] == "add_disclaimer":
            self["class"] = modify.AddDisclaimer
            self["headersonly"] = False
-
-            if "html_template" not in cfg and "html_file" in cfg:
-                cfg["html_template"] = cfg["html_file"]
-
-            if "text_template" not in cfg and "text_file" in cfg:
-                cfg["text_template"] = cfg["text_file"]
-
            if "error_policy" not in cfg:
                cfg["error_policy"] = "wrap"

            self.add_string_arg(
-                cfg, ("action", "html_template", "text_template",
-                      "error_policy"))
-            assert self["args"]["action"] in ("append", "prepend"), \
+                cfg, ["action", "html_template", "text_template",
+                      "error_policy"])
+            assert self["args"]["action"] in ["append", "prepend"], \
                f"{self['name']}: action: invalid value, " \
                f"should be 'append' or 'prepend'"
            assert self["args"]["error_policy"] in ("wrap",
@@ -132,22 +92,6 @@ class ActionConfig(BaseConfig):
                f"{self['name']}: error_policy: invalid value, " \
                f"should be 'wrap', 'ignore' or 'reject'"

-            try:
-                with open(self["args"]["html_template"], "r") as f:
-                    html = BeautifulSoup(f.read(), "html.parser")
-                    body = html.find('body')
-                    if body:
-                        # just use content within the body tag if present
-                        html = body
-                    self["args"]["html_template"] = html
-
-                with open(self["args"]["text_template"], "r") as f:
-                    self["args"]["text_template"] = f.read()
-
-            except IOError as e:
-                raise RuntimeError(
-                    f"{self['name']}: unable to open/read template file: {e}")
-
        elif self["type"] == "rewrite_links":
            self["class"] = modify.RewriteLinks
            self["headersonly"] = False
--- a/pymodmilter/modify.py
+++ b/pymodmilter/modify.py
@@ -20,6 +20,7 @@ __all__ = [
    "RewriteLinks"]

 import logging
+import re

 from base64 import b64encode
 from bs4 import BeautifulSoup
@@ -53,9 +54,17 @@ class AddHeader:
 class ModHeader:
    """Change the value of a mail header field."""
    def __init__(self, field, value, search=None):
-        self.field = field
        self.value = value
-        self.search = search
+
+        try:
+            self.field = re.compile(field, re.IGNORECASE)
+            if search is not None:
+                self.search = re.compile(
+                    search, re.MULTILINE + re.DOTALL + re.IGNORECASE)
+            else:
+                self.search = search
+        except re.error as e:
+            raise RuntimeError(e)

    def execute(self, milter, pretend=False,
                logger=logging.getLogger(__name__)):
@@ -101,8 +110,15 @@ class ModHeader:
 class DelHeader:
    """Delete a mail header field."""
    def __init__(self, field, value=None):
-        self.field = field
-        self.value = value
+        try:
+            self.field = re.compile(field, re.IGNORECASE)
+            if value is not None:
+                self.value = re.compile(
+                    value, re.MULTILINE + re.DOTALL + re.IGNORECASE)
+            else:
+                self.value = value
+        except re.error as e:
+            raise RuntimeError(e)

    def execute(self, milter, pretend=False,
                logger=logging.getLogger(__name__)):
@@ -196,8 +212,18 @@ def _wrap_message(milter):
 class AddDisclaimer:
    """Append or prepend a disclaimer to the mail body."""
    def __init__(self, text_template, html_template, action, error_policy):
-        self.text_template = text_template
-        self.html_template = html_template
+        try:
+            with open(text_template, "r") as f:
+                self.text_template = f.read()
+
+            with open(html_template, "r") as f:
+                html = BeautifulSoup(f.read(), "html.parser")
+
+        except IOError as e:
+            raise RuntimeError(e)
+
+        body = html.find('body')
+        self.html_template = body or html
        self.action = action
        self.error_policy = error_policy

--- a/pymodmilter/notify.py
+++ b/pymodmilter/notify.py
@@ -119,22 +119,26 @@ class EMailNotification(BaseNotification):
        self.mailfrom = envelope_from
        self.from_header = from_header
        self.subject = subject
-        self.template = open(template, "r").read()
-        self.embedded_imgs = []
-        for img_path in embed_imgs:
-            img = MIMEImage(open(img_path, "rb").read())
-            filename = basename(img_path)
-            img.add_header("Content-ID", f"<{filename}>")
-            self.embedded_imgs.append(img)
+        try:
+            self.template = open(template, "r").read()
+            self.embedded_imgs = []
+            for img_path in embed_imgs:
+                img = MIMEImage(open(img_path, "rb").read())
+                filename = basename(img_path)
+                img.add_header("Content-ID", f"<{filename}>")
+                self.embedded_imgs.append(img)

-        self.replacement_img = repl_img
-        self.strip_images = strip_imgs
+            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())
-            self.replacement_img.add_header(
-                "Content-ID", "<removed_for_security_reasons>")
+            if not strip_imgs and repl_img:
+                self.replacement_img = MIMEImage(
+                    open(repl_img, "rb").read())
+                self.replacement_img.add_header(
+                    "Content-ID", "<removed_for_security_reasons>")
+
+        except IOError as e:
+            raise RuntimeError(e)

        self.parser_lib = parser_lib

--- a/pymodmilter/run.py
+++ b/pymodmilter/run.py
@@ -79,7 +79,7 @@ def main():
        logger.setLevel(logging.INFO)

    try:
-        logger.debug("prepar milter configuration")
+        logger.debug("prepare milter configuration")
        cfg = ModifyMilterConfig(args.config, args.debug)

        if not args.debug:
@@ -106,6 +106,12 @@ def main():
        logger.error(e)
        sys.exit(255)

+    try:
+        ModifyMilter.set_config(cfg)
+    except (RuntimeError, ValueError) as e:
+        logger.error(e)
+        sys.exit(254)
+
    if args.test:
        print("Configuration OK")
        sys.exit(0)
@@ -123,7 +129,6 @@ def main():
    root_logger.addHandler(sysloghandler)

    logger.info("pymodmilter starting")
-    ModifyMilter.set_config(cfg)

    # register milter factory class
    Milter.factory = ModifyMilter