Added CLI, improved error handling, small bugfixes

.gitignore

@@ -0,0 +1,116 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

README.md

@@ -33,14 +33,14 @@ The following configuration options are mandatory in the global section:
 The following configuration options are mandatory in each quarantine section:
 * **regex**  
   Regular expression to filter e-mail headers.
-* **type**  
+* **quarantine_type**  
   One of the quarantine-types described below.
 * **action**  
   One of the actions described below.
-* **notification**  
+* **notification_type**  
   One of the notification types described below.
-* **whitelist**  
+* **whitelist_type**  
-  Database connection string (e.g. mysql://user:password@host:port) or NONE to disable whitelist.
+  One of the whitelist types described below.
 
 ### Quarantine types
 * **NONE**  
@@ -95,9 +95,18 @@ The following actions are available:
   Reject e-mails.
 
 
-### Whitelist
+### Whitelist types
-If a whitelist database connection string is configured, the following configuration options are mandatory:
+* **NONE**  
-* **whitelist_table**  
+  No whitelist will be used.
+
+* **DB**  
+  A database whitelist will be used. All database types supported by peewee are available.  
+
+  The following configuration options are mandatory for this whitelist type:
+* **whitelist_db_connection**  
+  Database connection string (e.g. mysql://user:password@host:port).  
+
+* **whitelist_db_table**  
   Database table to use.
 
 ## Developer information

pyquarantine/__init__.py

@@ -92,105 +92,127 @@ class QuarantineMilter(Milter.Base):
         return Milter.CONTINUE
 
     def eoh(self):
-        self.matched = None
+        try:
-        self.whitelist = whitelists.WhitelistCache()
+            self.matched = None
-        # iterate email headers
+            self.whitelist_cache = whitelists.WhitelistCache()
-        for header in self.headers:
+            # iterate email headers
-            self.logger.debug("{}: checking header '{}' against regex of every configured quarantine".format(self.queueid, header))
+            for header in self.headers:
-            # iterate quarantines
+                self.logger.debug("{}: checking header '{}' against regex of every configured quarantine".format(self.queueid, header))
-            for name, quarantine in self.config.items():
+                # iterate quarantines
-                if self.matched != None and quarantine["index"] == self.matched["index"]:
+                for name, quarantine in self.config.items():
-                    # a quarantine with higher precedence already matched, skip checks of quarantines with lower precedence
+                    if self.matched != None and quarantine["index"] == self.matched["index"]:
-                    self.logger.debug("{}: quarantine '{}' matched already, skip further checks of this header".format(self.queueid, name))
+                        # a quarantine with higher precedence already matched, skip checks of quarantines with lower precedence
+                        self.logger.debug("{}: quarantine '{}' matched already, skip further checks of this header".format(self.queueid, name))
+                        break
+                    self.logger.debug("{}: checking header against quarantine '{}'".format(self.queueid, name))
+                    # check if header matches regex
+                    if quarantine["regex_compiled"].match(header):
+                        whitelist = quarantine["whitelist_obj"]
+                        if whitelist != None:
+                            try:
+                                whitelisted_recipients = self.whitelist_cache.get_whitelisted_recipients(whitelist, self.mailfrom, self.recipients)
+                            except RuntimeError as e:
+                                self.logger.error("{}: unable to query whitelist: {}".format(self.queueid, e))
+                                return Milter.TEMPFAIL
+                            if len(whitelisted_recipients) == len(self.recipients):
+                                # all recipients are whitelisted, continue with header checks
+                                self.logger.debug("{}: header matched regex, but all recipients are whitelisted in quarantine '{}', continue checking this header".format(self.queueid, name))
+                                continue
+                        self.matched = quarantine
+                        # skip checks of this header with quarantines with lower precedence
+                        self.logger.debug("{}: header matched regex in quarantine '{}', further checks of this header will be skipped".format(self.queueid, name))
+                        break
+                if self.matched != None and self.matched["index"] == 0:
+                    self.logger.debug("{}: skipping checks of remaining headers, the quarantine with highest precedence matched already".format(self.queueid))
                     break
-                self.logger.debug("{}: checking header against quarantine '{}'".format(self.queueid, name))
+            if self.matched != None:
-                # check if header matches regex
+                self.logger.info("{}: email matched quarantine '{}'".format(self.queueid, self.matched["name"]))
-                if quarantine["regex_compiled"].match(header):
+                # one of the configured quarantines matched
-                    if quarantine["whitelist"] != None and \
+                if self.matched["quarantine_obj"] != None or self.matched["notification_obj"] != None:
-                            len(self.whitelist.get_whitelisted_recipients(quarantine["whitelist"], self.mailfrom, self.recipients)) == len(self.recipients):
+                    self.logger.debug("{}: initializing memory buffer to save email data".format(self.queueid))
-                        # all recipients are whitelisted, continue with header checks
+                    # quarantine or notification configured, initialize memory buffer to save mail
-                        self.logger.debug("{}: header matched regex, but all recipients are whitelisted in quarantine '{}', continue checking this header".format(self.queueid, name))
+                    self.fp = StringIO.StringIO()
-                        continue
+                    # write email headers to memory buffer
-                    self.matched = quarantine
+                    self.fp.write("{}\n".format("\n".join(self.headers)))
-                    # skip checks of this header with quarantines with lower precedence
+                else:
-                    self.logger.debug("{}: header matched regex in quarantine '{}', further checks of this header will be skipped".format(self.queueid, name))
+                    # quarantine and notification disabled, return configured action
-                    break
+                    self.logger.debug("{}: ".format(self.queueid))
-            if self.matched != None and self.matched["index"] == 0:
+                    self.logger.info("{}: quarantine and notification disabled, responding with configured action: {}".format(self.queueid, self.matched["action"].upper()))
-                self.logger.debug("{}: skipping checks of remaining headers, the quarantine with highest precedence matched already".format(self.queueid))
+                    return self.matched["milter_action"]
-                break
-        if self.matched != None:
-            self.logger.info("{}: email matched quarantine '{}'".format(self.queueid, self.matched["name"]))
-            # one of the configured quarantines matched
-            if self.matched["quarantine"] != None or self.matched["notification"] != None:
-                self.logger.debug("{}: initializing memory buffer to save email data".format(self.queueid))
-                # quarantine or notification configured, initialize memory buffer to save mail
-                self.fp = StringIO.StringIO()
-                # write email headers to memory buffer
-                self.fp.write("{}\n".format("\n".join(self.headers)))
             else:
-                # quarantine and notification disabled, return configured action
+                # no quarantine matched, accept mail
-                self.logger.debug("{}: ".format(self.queueid))
+                self.logger.info("{}: email passed clean".format(self.queueid))
-                self.logger.info("{}: quarantine and notification disabled, responding with configured action: {}".format(self.queueid, self.matched["action"].upper()))
+                return Milter.ACCEPT
-                return self.matched["milter_action"]
+            return Milter.CONTINUE
-        else:
+        except Exception as e:
-            # no quarantine matched, accept mail
+            self.logger.exception("an exception occured in eoh function: {}".format(e))
-            self.logger.info("{}: email passed clean".format(self.queueid))
+            return Milter.TEMPFAIL
-            return Milter.ACCEPT
-        return Milter.CONTINUE
 
     def body(self, chunk):
-        # save received body chunk
+        try:
-        self.fp.write(chunk)
+            # save received body chunk
+            self.fp.write(chunk)
+        except Exception as e:
+            self.logger.exception("an exception occured in body function: {}".format(e))
+            return Milter.TEMPFAIL
         return Milter.CONTINUE
 
     def eom(self):
-        if self.matched["whitelist"] != None:
+        try:
-            whitelisted_recipients = self.whitelist.get_whitelisted_recipients(self.matched["whitelist"], self.mailfrom, self.recipients)
+            if self.matched["whitelist_obj"] != None:
-            if len(whitelisted_recipients) > 0:
-                for recipient in whitelisted_recipients:
-                    self.recipients.remove(recipient)
-                self.fp.seek(0)
-                self.logger.info("{}: sending original email to whitelisted recipient(s): {}".format(self.queueid, "<{}>".format(">,<".join(whitelisted_recipients))))
                 try:
-                    mailer.sendmail(self.matched["smtp_host"], self.matched["smtp_port"], self.queueid, self.mailfrom, whitelisted_recipients, self.fp.read())
+                    whitelisted_recipients = self.whitelist_cache.get_whitelisted_recipients(self.matched["whitelist_obj"], self.mailfrom, self.recipients)
-                except Exception as e:
+                except RuntimeError as e:
-                    self.logger.error("{}: unable to send original email: {}".format(self.queueid, e))
+                    self.logger.error("{}: unable to query whitelist: {}".format(self.queueid, e))
                     return Milter.TEMPFAIL
-        if len(self.recipients) > 0:
+                if len(whitelisted_recipients) > 0:
-            quarantine_id = ""
+                    for recipient in whitelisted_recipients:
-            if self.matched["quarantine"] != None:
+                        self.recipients.remove(recipient)
-                # add email to quarantine
+                    self.fp.seek(0)
-                self.fp.seek(0)
+                    self.logger.info("{}: sending original email to whitelisted recipient(s): {}".format(self.queueid, "<{}>".format(">,<".join(whitelisted_recipients))))
-                try:
+                    try:
-                    quarantine_id = self.matched["quarantine"].add(self.queueid, self.mailfrom, self.recipients, fp=self.fp)
+                        mailer.sendmail(self.matched["smtp_host"], self.matched["smtp_port"], self.queueid, self.mailfrom, whitelisted_recipients, self.fp.read())
-                except Exception as e:
+                    except RuntimeError as e:
-                    self.logger.error("{}: unable to add email to quarantine: {}".format(self.queueid, e))
+                        self.logger.error("{}: unable to send original email: {}".format(self.queueid, e))
-                    return Milter.TEMPFAIL
+                        return Milter.TEMPFAIL
-                else:
+            if len(self.recipients) > 0:
-                    self.logger.info("{}: added email to quarantine of recipient(s): {}".format(self.queueid, "<{}>".format(">,<".join(self.recipients))))
+                quarantine_id = ""
-            if self.matched["notification"] != None:
+                if self.matched["quarantine_obj"] != None:
-                # notify
+                    # add email to quarantine
-                self.fp.seek(0)
+                    self.fp.seek(0)
-                try:
+                    self.logger.info("{}: adding email to quarantine of recipient(s): {}".format(self.queueid, "<{}>".format(">,<".join(self.recipients))))
-                    self.matched["notification"].notify(self.queueid, quarantine_id, self.subject, self.mailfrom, self.recipients, fp=self.fp)
+                    try:
-                except Exception as e:
+                        quarantine_id = self.matched["quarantine_obj"].add(self.queueid, self.mailfrom, self.recipients, fp=self.fp)
-                    self.logger.error("{}: unable to send notification(s): {}".format(self.queueid, e))
+                    except RuntimeError as e:
-                    return Milter.TEMPFAIL
+                        self.logger.error("{}: unable to add email to quarantine: {}".format(self.queueid, e))
-                else:
+                        return Milter.TEMPFAIL
-                    self.logger.info("{}: sent notification(s) to: {}".format(self.queueid, "<{}>".format(">,<".join(self.recipients))))
+                if self.matched["notification_obj"] != None:
-        self.fp.close()
+                    # notify
-        # return configured action
+                    self.fp.seek(0)
-        self.logger.info("{}: responding with configured action: {}".format(self.queueid, self.matched["action"].upper()))
+                    self.logger.info("{}: sending notification(s) to: {}".format(self.queueid, "<{}>".format(">,<".join(self.recipients))))
-        return self.matched["milter_action"]
+                    try:
+                        self.matched["notification_obj"].notify(self.queueid, quarantine_id, self.subject, self.mailfrom, self.recipients, fp=self.fp)
+                    except RuntimeError as e:
+                        self.logger.error("{}: unable to send notification(s): {}".format(self.queueid, e))
+                        return Milter.TEMPFAIL
+            self.fp.close()
+            # return configured action
+            self.logger.info("{}: responding with configured action: {}".format(self.queueid, self.matched["action"].upper()))
+            return self.matched["milter_action"]
+        except Exception as e:
+            self.logger.exception("an exception occured in eom function: {}".format(e))
+            return Milter.TEMPFAIL
 
 
-
+def generate_milter_config(configtest=False, config_files=[]):
-def generate_milter_config(configtest=False):
     "Generate the configuration for QuarantineMilter class."
     logger = logging.getLogger(__name__)
     # read config file
     parser = ConfigParser.ConfigParser()
-    config_files = parser.read(QuarantineMilter.get_configfiles())
+    if len(config_files) == 0:
+        config_files = parser.read(QuarantineMilter.get_configfiles())
+    else:
+        config_files = parser.read(config_files)
     if len(config_files) == 0:
         raise RuntimeError("config file not found")
     QuarantineMilter.set_configfiles(config_files)
@@ -219,7 +241,7 @@ def generate_milter_config(configtest=False):
         config[name] = dict(parser.items(name))
         config[name]["name"] = name
         # check if mandatory config options are present in config
-        for option in ["regex", "type", "notification", "action", "whitelist", "smtp_host", "smtp_port"]:
+        for option in ["regex", "quarantine_type", "notification_type", "action", "whitelist_type", "smtp_host", "smtp_port"]:
             if option not in config[name].keys() and \
                     option in config["global"].keys():
                 config[name][option] = config["global"][option]
@@ -234,7 +256,7 @@ def generate_milter_config(configtest=False):
         logger.debug("=> compiling regex '{}'".format(regex))
         config[name]["regex_compiled"] = re.compile(regex)
         # create quarantine instance
-        quarantine_type = config[name]["type"].lower()
+        quarantine_type = config[name]["quarantine_type"].lower()
         if quarantine_type in quarantines.quarantine_types.keys():
             logger.debug("=> initializing quarantine type '{}'".format(quarantine_type))
             quarantine = quarantines.quarantine_types[quarantine_type](name, config, configtest)
@@ -243,17 +265,18 @@ def generate_milter_config(configtest=False):
             quarantine = None
         else:
             raise RuntimeError("unknown quarantine_type '{}'".format(quarantine_type))
-        config[name]["quarantine"] = quarantine
+        config[name]["quarantine_obj"] = quarantine
         # create whitelist instance
-        whitelist = config[name]["whitelist"]
+        whitelist_type = config[name]["whitelist_type"].lower()
-        if whitelist.lower() == "none":
+        if whitelist_type in whitelists.whitelist_types.keys():
-            logger.debug("=> setting whitelist to NONE")
-            config[name]["whitelist"] = None
-        else:
             logger.debug("=> initializing whitelist database")
-            config[name]["whitelist"] = whitelists.Whitelist(name, config, configtest)
+            whitelist = whitelists.whitelist_types[whitelist_type](name, config, configtest)
+        else:
+            logger.debug("=> setting whitelist to NONE")
+            whitelist = None
+        config[name]["whitelist_obj"] = whitelist
         # create notification instance
-        notification_type = config[name]["notification"].lower()
+        notification_type = config[name]["notification_type"].lower()
         if notification_type in notifications.notification_types.keys():
             logger.debug("=> initializing notification type '{}'".format(notification_type))
             notification = notifications.notification_types[notification_type](name, config, configtest)
@@ -262,7 +285,7 @@ def generate_milter_config(configtest=False):
             notification = None
         else:
             raise RuntimeError("unknown notification type '{}'".format(notification_type))
-        config[name]["notification"] = notification
+        config[name]["notification_obj"] = notification
         # determining milter action for this quarantine
         action = config[name]["action"].upper()
         if action in QuarantineMilter.get_actions().keys():
@@ -272,8 +295,6 @@ def generate_milter_config(configtest=False):
             raise RuntimeError("unknown action '{}' configured for quarantine '{}'".format(action, name))
     # remove global section from config, every section should be a quarantine
     del(config["global"])
-    if configtest:
-        print("Configuration ok")
     return config
 
 

pyquarantine/cli.py

@@ -0,0 +1,346 @@
+#!/usr/bin/env python2
+#
+# PyQuarantine-Milter is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+# 
+# PyQuarantine-Milter is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the 
+# GNU General Public License for more details.
+# 
+# You should have received a copy of the GNU General Public License
+# along with PyQuarantineMilter.  If not, see <http://www.gnu.org/licenses/>.
+#
+
+import argparse
+import logging
+import logging.handlers
+import sys
+import time
+
+import pyquarantine
+
+
+
+def _get_quarantine_obj(config, quarantine):
+    if quarantine not in config.keys():
+        raise RuntimeError("invalid quarantine '{}'".format(quarantine))
+    return config[quarantine]["quarantine_obj"]
+
+
+
+def _get_whitelist_obj(config, quarantine):
+    if quarantine not in config.keys():
+        raise RuntimeError("invalid quarantine '{}'".format(quarantine))
+    return config[quarantine]["whitelist_obj"]
+
+
+
+def print_table(headers, keys, data):
+    if len(data) == 0:
+        return
+
+    # calculate length of each column
+    column_lengths = []
+    column_formats = []
+    for idx, header in enumerate(headers):
+        length = len(header)
+        key = keys[idx]
+        value_length=len((max(data.items(), key=lambda (k, v): len(v[key])))[1][key])
+        if value_length > length: length = value_length
+        column_lengths.append(length)
+        column_formats.append("{{:<{}}}".format(length))
+
+    # define row format
+    row_format = " | ".join(column_formats)
+
+    # define header/body separator
+    separators = []
+    for length in column_lengths:
+        separators.append("-" * length)
+    separator = "-+-".join(separators)
+
+    # print header and separator
+    print(row_format.format(*headers))
+    print(separator)
+
+    # print data
+    for key, value in data.items():
+        row = []
+        for entry in keys:
+            row.append(value[entry])
+        print(row_format.format(*row))
+
+
+
+def list_quarantines(config, args):
+    if args.batch:
+        print("\n".join(config.keys()))
+    else:
+        print_table(
+                ["Name", "Quarantine", "Notification", "Action"],
+                ["name", "quarantine_type", "notification_type", "action"],
+                config
+        )
+    return 0
+
+
+
+def list_quarantine_emails(config, args):
+    logger = logging.getLogger(__name__)
+
+    # get quarantine object
+    quarantine = _get_quarantine_obj(config, args.quarantine)
+    if quarantine == None:
+        raise RuntimeError("quarantine type is set to None, unable to list e-mails")
+
+    # find emails and transform some metadata values to strings
+    emails = quarantine.find(mailfrom=args.mailfrom, recipients=args.recipients, older_than=args.older_than)
+    for quarantine_id, metadata in emails.items():
+        emails[quarantine_id]["quarantine_id"] = quarantine_id
+        emails[quarantine_id]["recipient_str"] = ", ".join(metadata["recipients"])
+        emails[quarantine_id]["date_str"] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(metadata["date"]))
+
+    if args.batch:
+        print("\n".join(emails.keys()))
+    else:
+        if len(emails) == 0: logger.info("quarantine '{}' is empty".format(args.quarantine))
+        print_table(
+                ["Quarantine-ID", "From", "Recipient(s)", "Date"],
+                ["quarantine_id", "from", "recipient_str", "date_str"],
+                emails
+        )
+
+
+
+def list_whitelist(config, args):
+    logger = logging.getLogger(__name__)
+
+    # get whitelist object
+    whitelist = _get_whitelist_obj(config, args.quarantine)
+    if whitelist == None:
+        raise RuntimeError("whitelist type is set to None, unable to list entries")
+
+    # find whitelist entries
+    entries = whitelist.find(mailfrom=args.mailfrom, recipients=args.recipients, older_than=args.older_than)
+    if len(entries) == 0:
+        logger.info("whitelist of quarantine '{}' is empty".format(args.quarantine))
+    else:
+        # transform some values to strings
+        for entry_id, entry in entries.items():
+            entries[entry_id]["id"] = str(entry["id"])
+            entries[entry_id]["created_str"] = entry["created"].strftime('%Y-%m-%d %H:%M:%S')
+            entries[entry_id]["last_used_str"] = entry["last_used"].strftime('%Y-%m-%d %H:%M:%S')
+            entries[entry_id]["permanent_str"] = str(entry["permanent"])
+
+        print_table(
+                ["ID", "From", "To", "Created", "Last used", "Comment", "Permanent"],
+                ["id", "mailfrom", "recipient", "created_str", "last_used_str", "comment", "permanent_str"],
+                entries
+        )
+
+
+
+def add_whitelist_entry(config, args):
+    logger = logging.getLogger(__name__)
+
+    # get whitelist object
+    whitelist = _get_whitelist_obj(config, args.quarantine)
+    if whitelist == None:
+        raise RuntimeError("whitelist type is set to None, unable to add entries")
+
+    # check existing entries
+    entries = whitelist.check(args.mailfrom, args.recipient)
+    if len(entries) > 0:
+        # check if the exact entry exists already
+        for entry in entries.values():
+            if entry["mailfrom"] == args.mailfrom and entry["recipient"] == args.recipient:
+                raise RuntimeError("an entry with this from/to combination already exists")
+        if not args.force:
+            # the entry is already covered by others
+            for entry_id, entry in entries.items():
+                entries[entry_id]["created_str"] = entry["created"].strftime('%Y-%m-%d %H:%M:%S')
+                entries[entry_id]["last_used_str"] = entry["last_used"].strftime('%Y-%m-%d %H:%M:%S')
+                entries[entry_id]["permanent_str"] = str(entry["permanent"])
+            print_table(
+                    ["From", "To", "Created", "Last used", "Comment", "Permanent"],
+                    ["mailfrom", "recipient", "created_str", "last_used_str", "comment", "permanent_str"],
+                    entries
+            )
+            print("")
+            raise RuntimeError("from/to combination is already covered by the entries above, use --force to override.")
+    # add entry to whitelist
+    whitelist.add(args.mailfrom, args.recipient, args.comment, args.permanent)
+    logger.info("successfully added whitelist entry")
+
+
+
+def delete_whitelist_entry(config, args):
+    logger = logging.getLogger(__name__)
+    whitelist = _get_whitelist_obj(config, args.quarantine)
+    if whitelist == None:
+        raise RuntimeError("whitelist type is set to None, unable to delete entries")
+    whitelist.delete(args.whitelist_id)
+    logger.info("successfully deleted whitelist entry")
+
+
+
+def release_email(config, args):
+    logger = logging.getLogger(__name__)
+    quarantine = _get_quarantine_obj(config, args.quarantine)
+    if quarantine == None:
+        raise RuntimeError("quarantine type is set to None, unable to release e-mail")
+    quarantine.release(args.quarantine_id, args.recipient)
+    logger.info("successfully released e-mail [quarantine-id: {}] to '{}' from quarantine '{}'".format(args.quarantine_id, args.recipient, args.quarantine))
+
+
+
+def delete_email(config, args):
+    logger = logging.getLogger(__name__)
+    quarantine = _get_quarantine_obj(config, args.quarantine)
+    if quarantine == None:
+        raise RuntimeError("quarantine type is set to None, unable to delete e-mail")
+    quarantine.delete(args.quarantine_id, args.recipient)
+    logger.info("successfully deleted e-mail [quarantine-id: {}] to '{}' from quarantine '{}'".format(args.quarantine_id, args.recipient, args.quarantine))
+
+
+
+class StdErrFilter(logging.Filter):
+    def filter(self, rec):
+        return rec.levelno in (logging.ERROR, logging.WARNING)
+
+
+
+class StdOutFilter(logging.Filter):
+    def filter(self, rec):
+        return rec.levelno in (logging.DEBUG, logging.INFO)
+
+
+
+def main():
+    "PyQuarantine command-line interface."
+    # parse command line
+    formatter_class = lambda prog: argparse.HelpFormatter(prog, max_help_position=50, width=140)
+    parser = argparse.ArgumentParser(description="PyQuarantine CLI", formatter_class=formatter_class)
+    parser.add_argument("-c", "--config", help="Config files to read.", nargs="+", metavar="CFG",
+            default=pyquarantine.QuarantineMilter.get_configfiles())
+    parser.add_argument("-d", "--debug", help="Log debugging messages.", action="store_true")
+    parser.set_defaults(syslog=False)
+    subparsers = parser.add_subparsers()
+    # list command
+    list_parser = subparsers.add_parser("list", help="List available quarantines.", formatter_class=formatter_class)
+    list_parser.add_argument("-b", "--batch", help="Print results using only quarantine names, each on a new line.", action="store_true")
+    list_parser.set_defaults(func=list_quarantines)
+
+    # quarantine command group
+    quarantine_parser = subparsers.add_parser("quarantine", description="Manage quarantines.", help="Manage quarantines.", formatter_class=formatter_class)
+    quarantine_subparsers = quarantine_parser.add_subparsers()
+    # quarantine list command
+    quarantine_list_parser = quarantine_subparsers.add_parser("list", description="List e-mails in quarantines.", help="List e-mails in quarantine.", formatter_class=formatter_class)
+    quarantine_list_parser.add_argument("quarantine", metavar="QUARANTINE", help="Quarantine name.")
+    quarantine_list_parser.add_argument("-f", "--from", dest="mailfrom", help="Filter e-mails by from address.", default=None, nargs="+")
+    quarantine_list_parser.add_argument("-t", "--to", dest="recipients", help="Filter e-mails by recipient address.", default=None, nargs="+")
+    quarantine_list_parser.add_argument("-o", "--older-than", dest="older_than", help="Filter e-mails by age (days).", default=None, type=float)
+    quarantine_list_parser.add_argument("-b", "--batch", help="Print results using only e-mail quarantine IDs, each on a new line.", action="store_true")
+    quarantine_list_parser.set_defaults(func=list_quarantine_emails)
+    # quarantine release command
+    quarantine_release_parser = quarantine_subparsers.add_parser("release", description="Release e-mail from quarantine.", help="Release e-mail from quarantine.", formatter_class=formatter_class)
+    quarantine_release_parser.add_argument("quarantine", metavar="QUARANTINE", help="Quarantine name.")
+    quarantine_release_parser.add_argument("quarantine_id", metavar="ID", help="Quarantine ID.")
+    quarantine_release_parser.add_argument("-n", "--disable-syslog", dest="syslog", help="Disable syslog messages.", action="store_false")
+    quarantine_release_parser_group = quarantine_release_parser.add_mutually_exclusive_group(required=True)
+    quarantine_release_parser_group.add_argument("-t", "--to", dest="recipient", help="Release e-mail for one recipient address.")
+    quarantine_release_parser_group.add_argument("-a", "--all", help="Release e-mail for all recipients.", action="store_true")
+    quarantine_release_parser.set_defaults(func=release_email)
+    # quarantine delete command
+    quarantine_delete_parser = quarantine_subparsers.add_parser("delete", description="Delete e-mail from quarantine.", help="Delete e-mail from quarantine.", formatter_class=formatter_class)
+    quarantine_delete_parser.add_argument("quarantine", metavar="QUARANTINE", help="Quarantine name.")
+    quarantine_delete_parser.add_argument("quarantine_id", metavar="ID", help="Quarantine ID.")
+    quarantine_delete_parser.add_argument("-n", "--disable-syslog", dest="syslog", help="Disable syslog messages.", action="store_false")
+    quarantine_delete_parser_group = quarantine_delete_parser.add_mutually_exclusive_group(required=True)
+    quarantine_delete_parser_group.add_argument("-t", "--to", dest="recipient", help="Delete e-mail for one recipient address.")
+    quarantine_delete_parser_group.add_argument("-a", "--all", help="Delete e-mail for all recipients.", action="store_true")
+    quarantine_delete_parser.set_defaults(func=delete_email)
+
+    # whitelist command group
+    whitelist_parser = subparsers.add_parser("whitelist", description="Manage whitelists.", help="Manage whitelists.", formatter_class=formatter_class)
+    whitelist_subparsers = whitelist_parser.add_subparsers()
+    # whitelist list command
+    whitelist_list_parser = whitelist_subparsers.add_parser("list", description="List whitelist entries.", help="List whitelist entries.", formatter_class=formatter_class)
+    whitelist_list_parser.add_argument("quarantine", metavar="QUARANTINE", help="Quarantine name.")
+    whitelist_list_parser.add_argument("-f", "--from", dest="mailfrom", help="Filter entries by from address.", default=None, nargs="+")
+    whitelist_list_parser.add_argument("-t", "--to", dest="recipients", help="Filter entries by recipient address.", default=None, nargs="+")
+    whitelist_list_parser.add_argument("-o", "--older-than", dest="older_than", help="Filter e-mails by last used date (days).", default=None, type=float)
+    whitelist_list_parser.set_defaults(func=list_whitelist)
+    # whitelist add command
+    whitelist_add_parser = whitelist_subparsers.add_parser("add", description="Add whitelist entry.", help="Add whitelist entry.", formatter_class=formatter_class)
+    whitelist_add_parser.add_argument("quarantine", metavar="QUARANTINE", help="Quarantine name.")
+    whitelist_add_parser.add_argument("-f", "--from", dest="mailfrom", help="From address.", required=True)
+    whitelist_add_parser.add_argument("-t", "--to", dest="recipient", help="Recipient address.", required=True)
+    whitelist_add_parser.add_argument("-c", "--comment", help="Comment.", default="added by CLI", required=False)
+    whitelist_add_parser.add_argument("-p", "--permanent", help="Add a permanent entry.", action="store_true")
+    whitelist_add_parser.add_argument("--force", help="Force adding an entry, even if already covered by another entry.", action="store_true")
+    whitelist_add_parser.set_defaults(func=add_whitelist_entry)
+    # whitelist delete command
+    whitelist_delete_parser = whitelist_subparsers.add_parser("delete", description="Delete whitelist entry.", help="Delete whitelist entry.", formatter_class=formatter_class)
+    whitelist_delete_parser.add_argument("quarantine", metavar="QUARANTINE", help="Quarantine name.")
+    whitelist_delete_parser.add_argument("whitelist_id", metavar="ID", help="Whitelist ID.")
+    whitelist_delete_parser.set_defaults(func=delete_whitelist_entry)
+
+    args = parser.parse_args()
+
+    # setup logging
+    loglevel = logging.INFO
+    root_logger = logging.getLogger()
+    root_logger.setLevel(loglevel)
+
+    # setup console log
+    if args.debug:
+        formatter = logging.Formatter("%(levelname)s: [%(name)s] - %(message)s")
+    else:
+        formatter = logging.Formatter("%(levelname)s: %(message)s")
+    # stdout
+    stdouthandler = logging.StreamHandler(sys.stdout)
+    stdouthandler.setLevel(logging.DEBUG)
+    stdouthandler.setFormatter(formatter)
+    stdouthandler.addFilter(StdOutFilter())
+    root_logger.addHandler(stdouthandler)
+    # stderr
+    stderrhandler = logging.StreamHandler(sys.stderr)
+    stderrhandler.setLevel(logging.WARNING)
+    stderrhandler.setFormatter(formatter)
+    stderrhandler.addFilter(StdErrFilter())
+    root_logger.addHandler(stderrhandler)
+    logger = logging.getLogger(__name__)
+
+    # try to generate milter configs
+    try:
+        config = pyquarantine.generate_milter_config(config_files=args.config, configtest=True)
+    except RuntimeError as e:
+        logger.error(e)
+        sys.exit(255)
+
+    if args.syslog:
+        # setup syslog
+        sysloghandler = logging.handlers.SysLogHandler(address="/dev/log", facility=logging.handlers.SysLogHandler.LOG_MAIL)
+        sysloghandler.setLevel(loglevel)
+        if args.debug:
+            formatter = logging.Formatter("pyquarantine: [%(name)s] [%(levelname)s] %(message)s")
+        else:
+            formatter = logging.Formatter("pyquarantine: %(message)s")
+        sysloghandler.setFormatter(formatter)
+        root_logger.addHandler(sysloghandler)
+
+    # call the commands function
+    try:
+        args.func(config, args)
+    except RuntimeError as e:
+        logger.error(e)
+        sys.exit(1)
+
+
+
+if __name__ == "__main__":
+    main()

pyquarantine/mailer.py

@@ -26,6 +26,13 @@ process = None
 
 
 
+def smtp_send(smtp_host, smtp_port, mailfrom, recipient, mail):
+    s = smtplib.SMTP(host=smtp_host, port=smtp_port)
+    s.sendmail(mailfrom, [recipient], mail)
+    s.quit()
+
+
+
 def mailprocess():
     "Mailer process to send emails asynchronously."
     global logger
@@ -36,18 +43,17 @@ def mailprocess():
             if not m: break
             smtp_host, smtp_port, queueid, mailfrom, recipient, mail = m
             try:
-                s = smtplib.SMTP(host=smtp_host, port=smtp_port)
+                smtp_send(smtp_host, smtp_port, mailfrom, recipient, mail)
-                s.sendmail(mailfrom, [recipient], mail)
             except Exception as e:
                 logger.error("{}: error while sending email to <{}> via {}: {}".format(queueid, recipient, smtp_host, e))
             else:
                 logger.info("{}: email to <{}> sent successfully".format(queueid, recipient))
-                s.quit()
     except KeyboardInterrupt:
         pass
     logger.debug("mailer process terminated")
 
 
+
 def sendmail(smtp_host, smtp_port, queueid, mailfrom, recipients, mail):
     "Send an email."
     global logger
@@ -62,4 +68,7 @@ def sendmail(smtp_host, smtp_port, queueid, mailfrom, recipients, mail):
         logger.debug("starting mailer process")
         process.start()
     for recipient in recipients:
-        queue.put((smtp_host, smtp_port, queueid, mailfrom, recipient, mail))
+        try:
+            queue.put((smtp_host, smtp_port, queueid, mailfrom, recipient, mail), timeout=30)
+        except Queue.Full as e:
+            raise RuntimeError("e-mail queue is full")

pyquarantine/notifications.py

@@ -113,12 +113,12 @@ class EMailNotification(BaseNotification):
         self.subject = self.config["notification_email_subject"]
         try:
             self.template = open(self.config["notification_email_template"], "rb").read()
-        except Exception as e:
+        except IOError as e:
-            raise RuntimeError("error reading email template: {}".format(e))
+            raise RuntimeError("error reading template: {}".format(e))
         try:
             self.replacement_img = MIMEImage(open(self.config["notification_email_replacement_img"], "rb").read())
-        except Exception as e:
+        except IOError as e:
-            raise RuntimeError("error reading email replacement image: {}".format(e))
+            raise RuntimeError("error reading replacement image: {}".format(e))
         else:
             self.replacement_img.add_header("Content-ID", "<removed_for_security_reasons>")
 

pyquarantine/quarantines.py

@@ -14,14 +14,17 @@
 # along with PyQuarantineMilter.  If not, see <http://www.gnu.org/licenses/>.
 #
 
-import datetime
 import json
 import logging
 import os
 
+from calendar import timegm
+from datetime import datetime
+from glob import glob
 from shutil import copyfileobj
+from time import gmtime
 
-
+import mailer
 
 class BaseQuarantine(object):
     "Quarantine base class"
@@ -32,9 +35,25 @@ class BaseQuarantine(object):
         self.logger = logging.getLogger(__name__)
 
     def add(self, queueid, mailfrom, recipients, fp):
-        "Add mail to quarantine."
+        "Add e-mail to quarantine."
         return ""
 
+    def find(self, mailfrom=None, recipients=None, older_than=None):
+        "Find e-mails in quarantine."
+        return
+    
+    def get_metadata(self, quarantine_id):
+        "Return metadata of quarantined e-mail."
+        return
+
+    def delete(self, quarantine_id, recipient=None):
+        "Delete e-mail from quarantine."
+        return
+
+    def release(self, quarantine_id, recipient=None):
+        "Release e-mail from quarantine."
+        return
+
 
 
 class FileQuarantine(BaseQuarantine):
@@ -51,23 +70,139 @@ class FileQuarantine(BaseQuarantine):
         # check if quarantine directory exists and is writable
         if not os.path.isdir(self.directory) or not os.access(self.directory, os.W_OK):
             raise RuntimeError("file quarantine directory '{}' does not exist or is not writable".format(self.directory))
+        self._metadata_suffix = ".metadata"
+
+    def _save_datafile(self, quarantine_id, fp):
+        datafile = os.path.join(self.directory, quarantine_id)
+        try:
+            with open(datafile, "wb") as f:
+                copyfileobj(fp, f)
+        except IOError as e:
+            raise RuntimeError("unable save data file: {}".format(e))
+
+    def _save_metafile(self, quarantine_id, metadata):
+        metafile = os.path.join(self.directory, "{}{}".format(quarantine_id, self._metadata_suffix))
+        try:
+            with open(metafile, "wb") as f:
+                json.dump(metadata, f, indent=2)
+        except IOError as e:
+            raise RuntimeError("unable to save metadata file: {}".format(e))
+
+    def _remove(self, quarantine_id):
+        datafile = os.path.join(self.directory, quarantine_id)
+        metafile = "{}{}".format(datafile, self._metadata_suffix)
+        try:
+            os.remove(metafile)
+        except IOError as e:
+            raise RuntimeError("unable to remove metadata file: {}".format(e))
+        try:
+            os.remove(datafile)
+        except IOError as e:
+            raise RuntimeError("unable to remove data file: {}".format(e))
 
     def add(self, queueid, mailfrom, recipients, fp):
-        "Add mail to file quarantine and return quarantine-id."
+        "Add e-mail to file quarantine and return quarantine-id."
         super(FileQuarantine, self).add(queueid, mailfrom, recipients, fp)
-        quarantine_id = "{}_{}".format(datetime.datetime.now().strftime("%Y-%m-%d_%H%M%S"), queueid)
+        quarantine_id = "{}_{}".format(datetime.now().strftime("%Y%m%d%H%M%S"), queueid)
         # save mail
-        with open(os.path.join(self.directory, quarantine_id), "wb") as f:
+        self._save_datafile(quarantine_id, fp)
-            copyfileobj(fp, f)
         # save metadata
         metadata = {
             "from": mailfrom,
-            "recipients": recipients
+            "recipients": recipients,
+            "date": timegm(gmtime()),
+            "queue_id": queueid
         }
-        with open(os.path.join(self.directory, "{}.metadata".format(quarantine_id)), "wb") as f:
+        try:
-            json.dump(metadata, f, indent=2)
+            self._save_metafile(quarantine_id, metadata)
+        except RuntimeError as e:
+            datafile = os.path.join(self.directory, quarantine_id)
+            os.remove(datafile)
+            raise e
         return quarantine_id
 
+    def get_metadata(self, quarantine_id):
+        "Return metadata of quarantined e-mail."
+        super(FileQuarantine, self).get_metadata(quarantine_id)
+        metafile = os.path.join(self.directory, "{}{}".format(quarantine_id, self._metadata_suffix))
+        if not os.path.isfile(metafile):
+            raise RuntimeError("invalid quarantine id '{}'".format(quarantine_id))
+        try:
+            with open(metafile, "rb") as f:
+                metadata = json.load(f)
+        except IOError as e:
+            raise RuntimeError("unable to read metadata file: {}".format(e))
+        return metadata
+
+    def find(self, mailfrom=None, recipients=None, older_than=None):
+        "Find e-mails in quarantine."
+        super(FileQuarantine, self).find(mailfrom, recipients, older_than)
+        if type(mailfrom) == str: mailfrom = [mailfrom]
+        if type(recipients) == str: recipients = [recipients]
+        emails = {}
+        metafiles = glob(os.path.join(self.directory, "*{}".format(self._metadata_suffix)))
+        for metafile in metafiles:
+            if not os.path.isfile(metafile): continue
+            quarantine_id = os.path.basename(metafile[:-len(self._metadata_suffix)])
+            metadata = self.get_metadata(quarantine_id)
+            if older_than != None:
+                if timegm(gmtime()) - metadata["date"] < (older_than * 24 * 3600):
+                    continue
+            if mailfrom != None:
+                if metadata["from"] not in mailfrom:
+                    continue
+            if recipients != None:
+                if len(recipients) == 1 and recipients[0] not in metadata["recipients"]:
+                    continue
+                elif len(set(recipients + metadata["recipients"])) == len(recipients + metadata["recipients"]):
+                    continue
+            emails[quarantine_id] = metadata
+        return emails
+
+    def delete(self, quarantine_id, recipient=None):
+        "Delete e-mail in quarantine."
+        super(FileQuarantine, self).delete(quarantine_id, recipient)
+        try:
+            metadata = self.get_metadata(quarantine_id)
+        except RuntimeError as e:
+            raise RuntimeError("unable to delete e-mail: {}".format(e))
+        if recipient == None:
+            self._remove(quarantine_id)
+        else:
+            if recipient not in metadata["recipients"]:
+                raise RuntimeError("invalid recipient '{}'".format(recipient))
+            metadata["recipients"].remove(recipient)
+            if len(metadata["recipients"]) == 0:
+                self._remove(quarantine_id)
+            else:
+                self._save_metafile(quarantine_id, metadata)
+
+    def release(self, quarantine_id, recipient=None):
+        "Release e-mail from quarantine."
+        super(FileQuarantine, self).release(quarantine_id, recipient)
+        try:
+            metadata = self.get_metadata(quarantine_id)
+        except RuntimeError as e:
+            raise RuntimeError("unable to release e-mail: {}".format(e))
+        datafile = os.path.join(self.directory, quarantine_id)
+        if recipient != None:
+            if recipient not in metadata["recipients"]:
+                raise RuntimeError("invalid recipient '{}'".format(recipient))
+            recipients = [recipient]
+        else:
+            recipients = metadata["recipients"]
+        try:
+            with open(datafile, "rb") as f:
+                mail = f.read()
+        except IOError as e:
+            raise RuntimeError("unable to read data file: {}".format(e))
+        for recipient in recipients:
+            try:
+                mailer.smtp_send(self.config["smtp_host"], self.config["smtp_port"], metadata["from"], recipient, mail)
+            except Exception as e:
+                raise RuntimeError("error while sending e-mail to '{}': {}".format(recipient, e))
+            self.delete(quarantine_id, recipient)
+
 
 
 # list of quarantine types and their related quarantine classes

pyquarantine/run.py

@@ -27,13 +27,13 @@ import pyquarantine
 def main():
     "Run PyQuarantine-Milter."
     # parse command line
-    parser = argparse.ArgumentParser(description="Quarantine milter daemon",
+    parser = argparse.ArgumentParser(description="PyQuarantine milter daemon",
             formatter_class=lambda prog: argparse.HelpFormatter(prog, max_help_position=45, width=140))
-    parser.add_argument('-c', '--config', help='list of config files', nargs='+',
+    parser.add_argument("-c", "--config", help="List of config files to read.", nargs="+",
             default=pyquarantine.QuarantineMilter.get_configfiles())
-    parser.add_argument('-s', '--socket', help='socket used to communicatewith the MTA', required=True)
+    parser.add_argument("-s", "--socket", help="Socket used to communicate with the MTA.", required=True)
-    parser.add_argument('-d', '--debug', help='log debugging messages', action='store_true')
+    parser.add_argument("-d", "--debug", help="Log debugging messages.", action="store_true")
-    parser.add_argument('-t', '--test', help='check configuration', action='store_true')
+    parser.add_argument("-t", "--test", help="Check configuration.", action="store_true")
     args = parser.parse_args()
     # setup logging
     loglevel = logging.INFO
@@ -45,18 +45,19 @@ def main():
         syslog_name = "{}: [%(name)s] %(levelname)s".format(syslog_name)
     # set config files for milter class
     pyquarantine.QuarantineMilter.set_configfiles(args.config)
-    logger = logging.getLogger()
+    root_logger = logging.getLogger()
-    logger.setLevel(loglevel)
+    root_logger.setLevel(loglevel)
     # setup console log
     stdouthandler = logging.StreamHandler(sys.stdout)
     stdouthandler.setLevel(logging.DEBUG)
     formatter = logging.Formatter("%(message)s".format(logname))
     stdouthandler.setFormatter(formatter)
-    logger.addHandler(stdouthandler)
+    root_logger.addHandler(stdouthandler)
     logger = logging.getLogger(__name__)
     if args.test:
         try:
             pyquarantine.generate_milter_config(args.test)
+            print("Configuration ok")
         except RuntimeError as e:
             logger.error(e)
             sys.exit(255)
@@ -69,14 +70,15 @@ def main():
     sysloghandler.setLevel(loglevel)
     formatter = logging.Formatter("{}: %(message)s".format(syslog_name))
     sysloghandler.setFormatter(formatter)
-    logger.addHandler(sysloghandler)
+    root_logger.addHandler(sysloghandler)
     logger.info("PyQuarantine-Milter starting")
     try:
         # generate milter config
-        pyquarantine.QuarantineMilter.config = pyquarantine.generate_milter_config()
+        config = pyquarantine.generate_milter_config()
     except RuntimeError as e:
         logger.error(e)
         sys.exit(255)
+    pyquarantine.QuarantineMilter.config = config
     # register to have the Milter factory create instances of your class:
     Milter.factory = pyquarantine.QuarantineMilter
     Milter.set_exception_policy(Milter.TEMPFAIL)

pyquarantine/whitelists.py

@@ -23,6 +23,39 @@ from playhouse.db_url import connect
 
 
 
+class WhitelistBase(object):
+    "Whitelist base class"
+    def __init__(self, name, config, configtest=False):
+        self.name = name
+        self.config = config[name]
+        self.configtest = configtest
+        self.global_config = config["global"]
+        self.logger = logging.getLogger(__name__)
+        self.valid_entry_regex = re.compile(r"^[a-zA-Z0-9_.+-]*?(@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+)?$")
+
+    def check(self, mailfrom, recipient):
+        # check if mailfrom/recipient combination is whitelisted
+        return
+
+    def find(self, mailfrom=None, recipients=None, older_than=None):
+        "Find whitelist entries."
+        return
+
+    def add(self, mailfrom, recipient, comment, permanent):
+        "Add entry to whitelist."
+        # check if mailfrom and recipient are valid
+        if not self.valid_entry_regex.match(mailfrom):
+            raise RuntimeError("invalid from address")
+        if not self.valid_entry_regex.match(recipient):
+            raise RuntimeError("invalid recipient")
+        return
+
+    def delete(self, whitelist_id):
+        "Delete entry from whitelist."
+        return
+
+
+
 class WhitelistModel(peewee.Model):
     mailfrom = peewee.CharField()
     recipient = peewee.CharField()
@@ -40,45 +73,56 @@ class Meta(object):
 
 
 
-class Whitelist(object):
+class DatabaseWhitelist(WhitelistBase):
-    "Whitelist base class"
+    "Whitelist class to store whitelist in a database"
     _whitelists = {}
 
     def __init__(self, name, config, configtest=False):
-        self.name = name
+        super(DatabaseWhitelist, self).__init__(name, config, configtest)
-        self.config = config[name]
-        self.global_config = config["global"]
-        self.logger = logging.getLogger(__name__)
         # check if mandatory options are present in config
-        for option in ["whitelist_table"]:
+        for option in ["whitelist_db_connection", "whitelist_db_table"]:
             if option not in self.config.keys() and option in self.global_config.keys():
                 self.config[option] = self.global_config[option]
             if option not in self.config.keys():
                 raise RuntimeError("mandatory option '{}' not present in config section '{}' or 'global'".format(option, self.name))
-        self.tablename = self.config["whitelist_table"]
+        self.tablename = self.config["whitelist_db_table"]
-        connection_string = self.config["whitelist"]
+        connection_string = self.config["whitelist_db_connection"]
-        if connection_string in Whitelist._whitelists.keys():
+        if connection_string in DatabaseWhitelist._whitelists.keys():
-            self.db = Whitelist._whitelists[connection_string]
+            new_connection = False
-            return
+            self.db = DatabaseWhitelist._whitelists[connection_string]
-        try:
-            # connect to database
-            self.logger.debug("connecting to database '{}'".format(re.sub(r"(.*?://.*?):.*?(@.*)", r"\1:<PASSWORD>\2", connection_string)))
-            self.db = connect(connection_string)
-        except Exception as e:
-            raise RuntimeError("unable to connect to database: {}".format(e))
         else:
-            Whitelist._whitelists[connection_string] = self.db
+            new_connection = True
-        if configtest: return
+            try:
-        self.Meta = Meta
+                # connect to database
-        self.Meta.database = self.db
+                self.logger.debug("connecting to database '{}'".format(re.sub(r"(.*?://.*?):.*?(@.*)", r"\1:<PASSWORD>\2", connection_string)))
-        self.Meta.table_name = self.tablename
+                self.db = connect(connection_string)
-        self.Whitelist = type("WhitelistModel_{}".format(name), (WhitelistModel,), {
+            except Exception as e:
-            "Meta": self.Meta 
+                raise RuntimeError("unable to connect to database: {}".format(e))
+            DatabaseWhitelist._whitelists[connection_string] = self.db
+        self.meta = Meta
+        self.meta.database = self.db
+        self.meta.table_name = self.tablename
+        self.model = type("WhitelistModel_{}".format(name), (WhitelistModel,), {
+            "Meta": self.meta 
         })
-        try:
+        if new_connection and not self.configtest:
-            self.db.create_tables([self.Whitelist])
+            try:
-        except Exception as e:
+                self.db.create_tables([self.model])
-            raise RuntimeError("unable to initialize table '{}': {}".format(self.tablename, e))
+            except Exception as e:
+                raise RuntimeError("unable to initialize table '{}': {}".format(self.tablename, e))
+
+    def _entry_to_dict(self, entry):
+        result = {}
+        result[entry.id] = {
+            "id": entry.id,
+            "mailfrom": entry.mailfrom,
+            "recipient": entry.recipient,
+            "created": entry.created,
+            "last_used": entry.last_used,
+            "comment": entry.comment,
+            "permanent": entry.permanent
+        }
+        return result
 
     def get_weight(self, entry):
         value = 0
@@ -90,34 +134,78 @@ class Whitelist(object):
         return value
 
     def check(self, mailfrom, recipient):
+        # check if mailfrom/recipient combination is whitelisted
+        super(DatabaseWhitelist, self).check(mailfrom, recipient)
         # generate list of possible mailfroms
         self.logger.debug("query database for whitelist entries from <{}> to <{}>".format(mailfrom, recipient))
         mailfroms = [""]
-        if "@" in mailfrom:
+        if "@" in mailfrom and not mailfrom.startswith("@"):
             mailfroms.append("@{}".format(mailfrom.split("@")[1]))
         mailfroms.append(mailfrom)
         # generate list of possible recipients
         recipients = [""]
-        if "@" in recipient:
+        if "@" in recipient and not recipient.startswith("@"):
             recipients.append("@{}".format(recipient.split("@")[1]))
         recipients.append(recipient)
         # query the database
         try:
-            entries = list(self.Whitelist.select().where(self.Whitelist.mailfrom.in_(mailfroms), self.Whitelist.recipient.in_(recipients)))
+            entries = list(self.model.select().where(self.model.mailfrom.in_(mailfroms), self.model.recipient.in_(recipients)))
         except Exception as e:
-            entries = []
+            raise RuntimeError("unable to query database: {}".format(e))
-            self.logger.error("unable to query whitelist database: {}".format(e))
         if len(entries) == 0:
             # no whitelist entry found
-            return False 
+            return {}
-        print(entries)
         if len(entries) > 1:
-            entries.sort(key=lambda x: self.get_weight(x))
+            entries.sort(key=lambda x: self.get_weight(x), reverse=True)
         # use entry with the highest weight
-        entry = entries[-1]
+        entry = entries[0]
         entry.last_used = datetime.datetime.now()
         entry.save()
-        return True
+        result = {}
+        for entry in entries:
+            result.update(self._entry_to_dict(entry))
+        return result
+
+    def find(self, mailfrom=None, recipients=None, older_than=None):
+        "Find whitelist entries."
+        super(DatabaseWhitelist, self).find(mailfrom, recipients, older_than)
+        if type(mailfrom) == str: mailfrom = [mailfrom]
+        if type(recipients) == str: recipients = [recipients]
+        entries = {}
+        try:
+            for entry in list(self.model.select()):
+                if older_than != None: 
+                    if (datetime.datetime.now() - entry.last_used).total_seconds() < (older_than * 24 * 3600):
+                        continue
+                if mailfrom != None:
+                    if entry.mailfrom not in mailfrom:
+                        continue
+                if recipients != None:
+                    if entry.recipient not in recipients:
+                        continue
+                entries.update(self._entry_to_dict(entry))
+        except Exception as e:
+            raise RuntimeError("unable to query database: {}".format(e))
+        return entries
+
+    def add(self, mailfrom, recipient, comment, permanent):
+        "Add entry to whitelist."
+        super(DatabaseWhitelist, self).add(mailfrom, recipient, comment, permanent)
+        try:
+            self.model.create(mailfrom=mailfrom, recipient=recipient, comment=comment, permanent=permanent)
+        except Exception as e:
+            raise RuntimeError("unable to add entry to database: {}".format(e))
+
+    def delete(self, whitelist_id):
+        "Delete entry from whitelist."
+        super(DatabaseWhitelist, self).delete(whitelist_id)
+        try:
+            query = self.model.delete().where(self.model.id == whitelist_id)
+            deleted = query.execute()
+        except Exception as e:
+            raise RuntimeError("unable to delete entry from database: {}".format(e))
+        if deleted == 0:
+            raise RuntimeError("invalid whitelist id")
 
 
 
@@ -138,4 +226,9 @@ class WhitelistCache(object):
 
     def get_whitelisted_recipients(self, whitelist, mailfrom, recipients):
         self.load(whitelist, mailfrom, recipients)
-        return filter(lambda x: self.cache[whitelist][x] == True, self.cache[whitelist].keys())
+        return filter(lambda x: len(self.cache[whitelist][x]) > 0, self.cache[whitelist].keys())
+
+
+
+# list of whitelist types and their related whitelist classes
+whitelist_types = {"db": DatabaseWhitelist}

setup.py

@@ -27,9 +27,10 @@ setup(name = "pyquarantine",
     ],
     entry_points = {
         "console_scripts": [
-            "pyquarantine-milter=pyquarantine.run:main"
+            "pyquarantine-milter=pyquarantine.run:main",
+            "pyquarantine=pyquarantine.cli:main"
         ]
     },
     install_requires = ["pymilter", "peewee"],
     python_requires = ">=2.7,<3"
-) 
+)

test-pyquarantine

@@ -0,0 +1,9 @@
+#!/usr/bin/env python2
+
+import sys
+import pyquarantine.cli
+
+if __name__ == '__main__':
+    sys.exit(
+        pyquarantine.cli.main()
+    )

test-pyquarantine-milter

@@ -0,0 +1,9 @@
+#!/usr/bin/env python2
+
+import sys
+import pyquarantine.run
+
+if __name__ == '__main__':
+    sys.exit(
+        pyquarantine.run.main()
+    )