Added CLI, improved error handling, small bugfixes

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):
-    "Whitelist base class"
+class DatabaseWhitelist(WhitelistBase):
+    "Whitelist class to store whitelist in a database"
     _whitelists = {}
 
     def __init__(self, name, config, configtest=False):
-        self.name = name
-        self.config = config[name]
-        self.global_config = config["global"]
-        self.logger = logging.getLogger(__name__)
+        super(DatabaseWhitelist, self).__init__(name, config, configtest)
         # 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"]
-        connection_string = self.config["whitelist"]
-        if connection_string in Whitelist._whitelists.keys():
-            self.db = Whitelist._whitelists[connection_string]
-            return
-        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))
+        self.tablename = self.config["whitelist_db_table"]
+        connection_string = self.config["whitelist_db_connection"]
+        if connection_string in DatabaseWhitelist._whitelists.keys():
+            new_connection = False
+            self.db = DatabaseWhitelist._whitelists[connection_string]
         else:
-            Whitelist._whitelists[connection_string] = self.db
-        if configtest: return
-        self.Meta = Meta
-        self.Meta.database = self.db
-        self.Meta.table_name = self.tablename
-        self.Whitelist = type("WhitelistModel_{}".format(name), (WhitelistModel,), {
-            "Meta": self.Meta 
+            new_connection = True
+            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))
+            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:
-            self.db.create_tables([self.Whitelist])
-        except Exception as e:
-            raise RuntimeError("unable to initialize table '{}': {}".format(self.tablename, e))
+        if new_connection and not self.configtest:
+            try:
+                self.db.create_tables([self.model])
+            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 = []
-            self.logger.error("unable to query whitelist database: {}".format(e))
+            raise RuntimeError("unable to query database: {}".format(e))
         if len(entries) == 0:
             # no whitelist entry found
-            return False 
-        print(entries)
+            return {}
         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}