spacefreak/pyquarantine-milter
1
0
Fork 0
Code Issues Pull Requests Releases Activity

Make source PEP8 conform

Browse Source
This commit is contained in:
Thomas Oettli
2019-10-17 22:25:10 +02:00
parent 6ea167bc52
commit 89a01d92c8
7 changed files with 865 additions and 366 deletions
Download Patch File Download Diff File Expand all files Collapse all files

102
pyquarantine/whitelists.py
View File

@@ -2,34 +2,36 @@
# 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
# 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 datetime
import logging
import peewee
import re
import sys
from datetime import datetime
from playhouse.db_url import connect
class WhitelistBase(object):
"Whitelist base class"
def __init__(self, global_config, config, configtest=False):
self.global_config = global_config
self.config = config
self.configtest = configtest
self.name = config["name"]
self.logger = logging.getLogger(__name__)
self.valid_entry_regex = re.compile(r"^[a-zA-Z0-9_.+-]*?(@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+)?$")
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."
@@ -56,15 +58,16 @@ class WhitelistBase(object):
class WhitelistModel(peewee.Model):
mailfrom = peewee.CharField()
recipient = peewee.CharField()
created = peewee.DateTimeField(default=datetime.datetime.now)
last_used = peewee.DateTimeField(default=datetime.datetime.now)
created = peewee.DateTimeField(default=datetime.now)
last_used = peewee.DateTimeField(default=datetime.now)
comment = peewee.TextField(default="")
permanent = peewee.BooleanField(default=False)
class Meta(object):
indexes = (
(('mailfrom', 'recipient'), True), # trailing comma is mandatory if only one index should be created
# trailing comma is mandatory if only one index should be created
(('mailfrom', 'recipient'), True),
)
@@ -74,14 +77,21 @@ class DatabaseWhitelist(WhitelistBase):
_db_tables = {}
def __init__(self, global_config, config, configtest=False):
super(DatabaseWhitelist, self).__init__(global_config, config, configtest)
super(
DatabaseWhitelist,
self).__init__(
global_config,
config,
configtest)
# check if mandatory options are present in config
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))
raise RuntimeError(
"mandatory option '{}' not present in config section '{}' or 'global'".format(
option, self.name))
tablename = self.config["whitelist_db_table"]
connection_string = self.config["whitelist_db_connection"]
@@ -91,10 +101,16 @@ class DatabaseWhitelist(WhitelistBase):
else:
try:
# connect to database
self.logger.debug("connecting to database '{}'".format(re.sub(r"(.*?://.*?):.*?(@.*)", r"\1:<PASSWORD>\2", connection_string)))
self.logger.debug(
"connecting to database '{}'".format(
re.sub(
r"(.*?://.*?):.*?(@.*)",
r"\1:<PASSWORD>\2",
connection_string)))
db = connect(connection_string)
except Exception as e:
raise RuntimeError("unable to connect to database: {}".format(e))
raise RuntimeError(
"unable to connect to database: {}".format(e))
DatabaseWhitelist._db_connections[connection_string] = db
@@ -103,7 +119,7 @@ class DatabaseWhitelist(WhitelistBase):
self.meta.database = db
self.meta.table_name = tablename
self.model = type("WhitelistModel_{}".format(self.name), (WhitelistModel,), {
"Meta": self.meta
"Meta": self.meta
})
if connection_string not in DatabaseWhitelist._db_tables.keys():
@@ -115,7 +131,9 @@ class DatabaseWhitelist(WhitelistBase):
try:
db.create_tables([self.model])
except Exception as e:
raise RuntimeError("unable to initialize table '{}': {}".format(tablename, e))
raise RuntimeError(
"unable to initialize table '{}': {}".format(
tablename, e))
def _entry_to_dict(self, entry):
result = {}
@@ -144,7 +162,9 @@ class DatabaseWhitelist(WhitelistBase):
super(DatabaseWhitelist, self).check(mailfrom, recipient)
# generate list of possible mailfroms
self.logger.debug("query database for whitelist entries from <{}> to <{}>".format(mailfrom, recipient))
self.logger.debug(
"query database for whitelist entries from <{}> to <{}>".format(
mailfrom, recipient))
mailfroms = [""]
if "@" in mailfrom and not mailfrom.startswith("@"):
mailfroms.append("@{}".format(mailfrom.split("@")[1]))
@@ -158,7 +178,10 @@ class DatabaseWhitelist(WhitelistBase):
# query the database
try:
entries = list(self.model.select().where(self.model.mailfrom.in_(mailfroms), self.model.recipient.in_(recipients)))
entries = list(
self.model.select().where(
self.model.mailfrom.in_(mailfroms),
self.model.recipient.in_(recipients)))
except Exception as e:
raise RuntimeError("unable to query database: {}".format(e))
@@ -171,7 +194,7 @@ class DatabaseWhitelist(WhitelistBase):
# use entry with the highest weight
entry = entries[0]
entry.last_used = datetime.datetime.now()
entry.last_used = datetime.now()
entry.save()
result = {}
for entry in entries:
@@ -183,21 +206,23 @@ class DatabaseWhitelist(WhitelistBase):
"Find whitelist entries."
super(DatabaseWhitelist, self).find(mailfrom, recipients, older_than)
if type(mailfrom) == str: mailfrom = [mailfrom]
if type(recipients) == str: recipients = [recipients]
if isinstance(mailfrom, str):
mailfrom = [mailfrom]
if isinstance(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):
if older_than is not None:
if (datetime.now() - entry.last_used).total_seconds() < (older_than * 86400):
continue
if mailfrom != None:
if mailfrom is not None:
if entry.mailfrom not in mailfrom:
continue
if recipients != None:
if recipients is not None:
if entry.recipient not in recipients:
continue
@@ -209,10 +234,20 @@ class DatabaseWhitelist(WhitelistBase):
def add(self, mailfrom, recipient, comment, permanent):
"Add entry to whitelist."
super(DatabaseWhitelist, self).add(mailfrom, recipient, comment, permanent)
super(
DatabaseWhitelist,
self).add(
mailfrom,
recipient,
comment,
permanent)
try:
self.model.create(mailfrom=mailfrom, recipient=recipient, comment=comment, permanent=permanent)
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))
@@ -224,7 +259,8 @@ class DatabaseWhitelist(WhitelistBase):
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))
raise RuntimeError(
"unable to delete entry from database: {}".format(e))
if deleted == 0:
raise RuntimeError("invalid whitelist id")
@@ -239,15 +275,19 @@ class WhitelistCache(object):
self.check(whitelist, mailfrom, recipient)
def check(self, whitelist, mailfrom, recipient):
if whitelist not in self.cache.keys(): self.cache[whitelist] = {}
if recipient not in self.cache[whitelist].keys(): self.cache[whitelist][recipient] = None
if self.cache[whitelist][recipient] == None:
self.cache[whitelist][recipient] = whitelist.check(mailfrom, recipient)
if whitelist not in self.cache.keys():
self.cache[whitelist] = {}
if recipient not in self.cache[whitelist].keys():
self.cache[whitelist][recipient] = None
if self.cache[whitelist][recipient] is None:
self.cache[whitelist][recipient] = whitelist.check(
mailfrom, recipient)
return self.cache[whitelist][recipient]
def get_whitelisted_recipients(self, whitelist, mailfrom, recipients):
self.load(whitelist, mailfrom, recipients)
return list(filter(lambda x: self.cache[whitelist][x], self.cache[whitelist].keys()))
return list(
filter(lambda x: self.cache[whitelist][x], self.cache[whitelist].keys()))
# list of whitelist types and their related whitelist classes