From 01ae131088e10d0b36628f286f4c0ea03581a1e5 Mon Sep 17 00:00:00 2001 From: Thomas Oettli Date: Thu, 30 Sep 2021 23:40:50 +0200 Subject: [PATCH] adapt all CLI functions to new code structure --- pyquarantine/__init__.py | 7 -- pyquarantine/cli.py | 64 +++++--------- pyquarantine/conditions.py | 3 + pyquarantine/config.py | 22 +++-- pyquarantine/notify.py | 5 +- pyquarantine/rule.py | 9 +- pyquarantine/storage.py | 169 +++++++++++++++++++++++++++---------- 7 files changed, 170 insertions(+), 109 deletions(-) diff --git a/pyquarantine/__init__.py b/pyquarantine/__init__.py index c46681b..7472c63 100644 --- a/pyquarantine/__init__.py +++ b/pyquarantine/__init__.py @@ -71,14 +71,7 @@ class ModifyMilter(Milter.Base): logger = logging.getLogger(__name__) logger.setLevel(ModifyMilter._loglevel) for idx, rule_cfg in enumerate(cfg["rules"]): - if "name" not in rule_cfg: - rule_cfg["name"] = f"rule#{idx}" - if "loglevel" not in rule_cfg: - rule_cfg["loglevel"] = cfg["loglevel"] - if "pretend" not in rule_cfg: - rule_cfg["pretend"] = cfg["pretend"] rule = Rule(rule_cfg, local_addrs, debug) - logger.debug(rule) ModifyMilter._rules.append(rule) diff --git a/pyquarantine/cli.py b/pyquarantine/cli.py index d7d733e..28f6029 100644 --- a/pyquarantine/cli.py +++ b/pyquarantine/cli.py @@ -21,6 +21,7 @@ import sys import time from pyquarantine.config import get_milter_config +from pyquarantine.storage import Quarantine from pyquarantine import __version__ as version @@ -28,22 +29,12 @@ def _get_quarantine(quarantines, name): try: quarantine = next((q for q in quarantines if q.name == name)) except StopIteration: - raise RuntimeError("invalid quarantine 'name'") + raise RuntimeError(f"invalid quarantine '{name}'") return quarantine -def _get_storage(quarantines, name): - quarantine = _get_quarantine(quarantines, name) - storage = quarantine.get_storage() - if not storage: - raise RuntimeError( - "storage type is set to NONE") - return storage - - def _get_notification(quarantines, name): - quarantine = _get_quarantine(quarantines, name) - notification = quarantine.get_notification() + notification = _get_quarantine(quarantines, name).notification if not notification: raise RuntimeError( "notification type is set to NONE") @@ -51,8 +42,7 @@ def _get_notification(quarantines, name): def _get_whitelist(quarantines, name): - quarantine = _get_quarantine(quarantines, name) - whitelist = quarantine.get_whitelist() + whitelist = _get_quarantine(quarantines, name).whitelist if not whitelist: raise RuntimeError( "whitelist type is set to NONE") @@ -106,21 +96,15 @@ def list_quarantines(quarantines, args): else: qlist = [] for q in quarantines: - storage = q.get_storage() - if storage: - storage_type = q.get_storage().storage_type - else: - storage_type = "NONE" + storage_type = type(q.storage).__name__ - notification = q.get_notification() - if notification: - notification_type = q.get_notification().notification_type + if q.notification: + notification_type = type(q.notification).__name__ else: notification_type = "NONE" - whitelist = q.get_whitelist() - if whitelist: - whitelist_type = q.get_whitelist().whitelist_type + if q.whitelist: + whitelist_type = type(q.whitelist).__name__ else: whitelist_type = "NONE" @@ -129,7 +113,7 @@ def list_quarantines(quarantines, args): "storage": storage_type, "notification": notification_type, "whitelist": whitelist_type, - "action": q.action}) + "action": q.milter_action}) print_table( [("Name", "name"), ("Storage", "storage"), @@ -142,7 +126,8 @@ def list_quarantines(quarantines, args): def list_quarantine_emails(quarantines, args): logger = logging.getLogger(__name__) - storage = _get_storage(quarantines, args.quarantine) + storage = _get_quarantine(quarantines, args.quarantine).storage + # find emails and transform some metadata values to strings rows = [] emails = storage.find( @@ -156,9 +141,9 @@ def list_quarantine_emails(quarantines, args): metadata["date"])) row["mailfrom"] = metadata["mailfrom"] row["recipient"] = metadata["recipients"].pop(0) - if "subject" not in emails[storage_id]["headers"].keys(): - emails[storage_id]["headers"]["subject"] = "" - row["subject"] = emails[storage_id]["headers"]["subject"][:60].strip() + if "subject" not in emails[storage_id]: + emails[storage_id]["subject"] = "" + row["subject"] = emails[storage_id]["subject"][:60].strip() rows.append(row) if metadata["recipients"]: @@ -223,7 +208,7 @@ def add_whitelist_entry(quarantines, args): whitelist = _get_whitelist(quarantines, args.quarantine) # check existing entries - entries = whitelist.check(args.mailfrom, args.recipient) + entries = whitelist.check(args.mailfrom, args.recipient, logger) if entries: # check if the exact entry exists already for entry in entries.values(): @@ -281,16 +266,15 @@ def release(quarantines, args): def delete(quarantines, args): logger = logging.getLogger(__name__) - storage = _get_storage(quarantines, args.quarantine) + storage = _get_quarantine(quarantines, args.quarantine).storage storage.delete(args.quarantine_id, args.recipient) logger.info("quarantined email deleted successfully") def get(quarantines, args): - storage = _get_storage(quarantines, args.quarantine) - fp, _ = storage.get_mail(args.quarantine_id) - print(fp.read().decode()) - fp.close() + storage = _get_quarantine(quarantines, args.quarantine).storage + _, msg = storage.get_mail(args.quarantine_id) + print(msg.as_string()) class StdErrFilter(logging.Filter): @@ -597,10 +581,8 @@ def main(): for rule in cfg["rules"]: for action in rule["actions"]: if action["type"] == "quarantine": - quarantines.append(action) - - print(quarantines) - sys.exit(0) + quarantines.append( + Quarantine(action, [], args.debug)) if args.syslog: # setup syslog @@ -618,7 +600,7 @@ def main(): # call the commands function try: - args.func(cfg, args) + args.func(quarantines, args) except RuntimeError as e: logger.error(e) sys.exit(1) diff --git a/pyquarantine/conditions.py b/pyquarantine/conditions.py index 3ec4b88..bfd2a8f 100644 --- a/pyquarantine/conditions.py +++ b/pyquarantine/conditions.py @@ -80,6 +80,9 @@ class Conditions: cfg.append(f"whitelist={self.whitelist}") return "Conditions(" + ", ".join(cfg) + ")" + def get_whitelist(self): + return self.whitelist + def match_host(self, host): logger = CustomLogger( self.logger, {"name": self.cfg["name"]}) diff --git a/pyquarantine/config.py b/pyquarantine/config.py index c729ad3..5242593 100644 --- a/pyquarantine/config.py +++ b/pyquarantine/config.py @@ -230,7 +230,7 @@ class NotifyConfig(BaseConfig): class QuarantineConfig(BaseConfig): JSON_SCHEMA = { "type": "object", - "required": ["store"], + "required": ["store", "smtp_host", "smtp_port"], "additionalProperties": False, "properties": { "name": {"type": "string"}, @@ -238,7 +238,9 @@ class QuarantineConfig(BaseConfig): "milter_action": {"type": "string"}, "reject_reason": {"type": "string"}, "whitelist": {"type": "object"}, - "store": {"type": "object"}}} + "store": {"type": "object"}, + "smtp_host": {"type": "string"}, + "smtp_port": {"type": "number"}}} def __init__(self, config, rec=True): super().__init__(config) @@ -264,10 +266,10 @@ class ActionConfig(BaseConfig): JSON_SCHEMA = { "type": "object", - "required": ["type", "args"], + "required": ["name", "type", "args"], "additionalProperties": False, "properties": { - "name": {"type": "string", "default": "action"}, + "name": {"type": "string"}, "loglevel": {"type": "string", "default": "info"}, "pretend": {"type": "boolean", "default": False}, "conditions": {"type": "object"}, @@ -285,10 +287,10 @@ class ActionConfig(BaseConfig): class RuleConfig(BaseConfig): JSON_SCHEMA = { "type": "object", - "required": ["actions"], + "required": ["name", "actions"], "additionalProperties": False, "properties": { - "name": {"type": "string", "default": "rule"}, + "name": {"type": "string"}, "loglevel": {"type": "string", "default": "info"}, "pretend": {"type": "boolean", "default": False}, "conditions": {"type": "object"}, @@ -302,6 +304,10 @@ class RuleConfig(BaseConfig): actions = [] for idx, action in enumerate(self["actions"]): + if "loglevel" not in action: + action["loglevel"] = config["loglevel"] + if "pretend" not in action: + action["pretend"] = config["pretend"] actions.append(ActionConfig(action, rec)) self["actions"] = actions @@ -331,6 +337,10 @@ class MilterConfig(BaseConfig): if rec: rules = [] for idx, rule in enumerate(self["rules"]): + if "loglevel" not in rule: + rule["loglevel"] = config["loglevel"] + if "pretend" not in rule: + rule["pretend"] = config["pretend"] rules.append(RuleConfig(rule, rec)) self["rules"] = rules diff --git a/pyquarantine/notify.py b/pyquarantine/notify.py index 56ca4ba..9eb35ba 100644 --- a/pyquarantine/notify.py +++ b/pyquarantine/notify.py @@ -225,7 +225,7 @@ class EMailNotification(BaseNotification): return soup def notify(self, msg, qid, mailfrom, recipients, logger, - template_vars=defaultdict(str), synchronous=False): + template_vars={}, synchronous=False): "Notify recipients via email." # extract body from email soup = self.get_email_body_soup(msg, logger) @@ -336,6 +336,9 @@ class Notify: class_name = type(self._notification).__name__ return f"{class_name}(" + ", ".join(cfg) + ")" + def get_notification(self): + return self._notification + def execute(self, milter): logger = CustomLogger( self.logger, {"name": self.cfg["name"], "qid": milter.qid}) diff --git a/pyquarantine/rule.py b/pyquarantine/rule.py index e0e003b..35cc54d 100644 --- a/pyquarantine/rule.py +++ b/pyquarantine/rule.py @@ -33,14 +33,7 @@ class Rule: self.actions = [] for idx, action_cfg in enumerate(cfg["actions"]): - if "name" in action_cfg: - action_cfg["name"] = f"{cfg['name']}: {action_cfg['name']}" - else: - action_cfg["name"] = f"action#{idx}" - if "loglevel" not in action_cfg: - action_cfg["loglevel"] = cfg["loglevel"] - if "pretend" not in action_cfg: - action_cfg["pretend"] = cfg["pretend"] + action_cfg["name"] = f"{cfg['name']}: {action_cfg['name']}" self.actions.append(Action(action_cfg, local_addrs, debug)) def __str__(self): diff --git a/pyquarantine/storage.py b/pyquarantine/storage.py index 597e2dd..d1b333c 100644 --- a/pyquarantine/storage.py +++ b/pyquarantine/storage.py @@ -24,10 +24,13 @@ import os from calendar import timegm from datetime import datetime +from email import message_from_binary_file +from email.policy import SMTPUTF8 from glob import glob from time import gmtime -from pyquarantine.base import CustomLogger +from pyquarantine import mailer +from pyquarantine.base import CustomLogger, MilterMessage from pyquarantine.conditions import Conditions from pyquarantine.config import ActionConfig from pyquarantine.notify import Notify @@ -44,7 +47,7 @@ class BaseMailStorage: self.metavar = metavar self.pretend = False - def add(self, data, qid, mailfrom="", recipients=[]): + def add(self, data, qid, mailfrom, recipients, subject, variables): "Add email to storage." return ("", "") @@ -152,15 +155,25 @@ class FileMailStorage(BaseMailStorage): except IOError as e: raise RuntimeError(f"unable to remove file: {e}") - def add(self, data, qid, mailfrom="", recipients=[], subject=""): + def add(self, data, qid, mailfrom, recipients, subject, variables, logger): "Add email to file storage and return storage id." - super().add(data, qid, mailfrom, recipients) + super().add(data, qid, mailfrom, recipients, subject, variables) storage_id = self.get_storageid(qid) metafile, datafile = self._get_file_paths(storage_id) + if self.metavar: + variables[f"{self.metavar}_ID"] = storage_id + variables[f"{self.metavar}_DATAFILE"] = datafile + if self.metadata: + variables[f"{self.metavar}_METAFILE"] = metafile + + if self.pretend: + return + # save mail self._save_datafile(datafile, data) + logger.info(f"stored message in file {datafile}") if not self.metadata: return storage_id, None, datafile @@ -169,9 +182,11 @@ class FileMailStorage(BaseMailStorage): metadata = { "mailfrom": mailfrom, "recipients": recipients, + "date": timegm(gmtime()), "subject": subject, "timestamp": timegm(gmtime()), - "queue_id": qid} + "queue_id": qid, + "vars": variables} try: self._save_metafile(metafile, metadata) @@ -179,8 +194,6 @@ class FileMailStorage(BaseMailStorage): os.remove(datafile) raise e - return storage_id, metafile, datafile - def execute(self, milter, logger): if self.original: milter.fp.seek(0) @@ -194,19 +207,8 @@ class FileMailStorage(BaseMailStorage): recipients = list(milter.msginfo["rcpts"]) subject = milter.msg["subject"] or "" - if not self.pretend: - storage_id, metafile, datafile = self.add( - data(), milter.qid, mailfrom, recipients, subject) - logger.info(f"stored message in file {datafile}") - else: - storage_id = self.get_storageid(milter.qid) - metafile, datafile = self._get_file_paths(storage_id) - - if self.metavar: - milter.msginfo["vars"][f"{self.metavar}_ID"] = storage_id - milter.msginfo["vars"][f"{self.metavar}_DATAFILE"] = datafile - if self.metadata: - milter.msginfo["vars"][f"{self.metavar}_METAFILE"] = metafile + self.add(data(), milter.qid, mailfrom, recipients, subject, + milter.msginfo["vars"], logger) def get_metadata(self, storage_id): "Return metadata of email in storage." @@ -275,7 +277,6 @@ class FileMailStorage(BaseMailStorage): def delete(self, storage_id, recipients=None): "Delete email from storage." super().delete(storage_id, recipients) - if not recipients or not self.metadata: self._remove(storage_id) return @@ -305,10 +306,13 @@ class FileMailStorage(BaseMailStorage): metadata = self.get_metadata(storage_id) _, datafile = self._get_file_paths(storage_id) try: - data = open(datafile, "rb").read() + with open(datafile, "rb") as fh: + msg = message_from_binary_file( + fh, _class=MilterMessage, policy=SMTPUTF8.clone( + refold_source='none')) except IOError as e: raise RuntimeError(f"unable to open email data file: {e}") - return (metadata, data) + return (metadata, msg) class Store: @@ -334,6 +338,9 @@ class Store: class_name = type(self._storage).__name__ return f"{class_name}(" + ", ".join(cfg) + ")" + def get_storage(self): + return self._storage + def execute(self, milter): logger = CustomLogger( self.logger, {"name": self.cfg["name"], "qid": milter.qid}) @@ -349,16 +356,19 @@ class Quarantine: self.logger = logging.getLogger(cfg["name"]) self.logger.setLevel(cfg.get_loglevel(debug)) - store_cfg = ActionConfig({ + storage_cfg = ActionConfig({ "name": cfg["name"], "loglevel": cfg["loglevel"], "pretend": cfg["pretend"], "type": "store", "args": cfg["args"]["store"].get_config()}) - store_cfg["args"]["metadata"] = True - self.store = Store(store_cfg, local_addrs, debug) + storage_cfg["args"]["metadata"] = True + self._storage = Store(storage_cfg, local_addrs, debug) - self.notify = None + self.smtp_host = cfg["args"]["smtp_host"] + self.smtp_port = cfg["args"]["smtp_port"] + + self._notification = None if "notify" in cfg["args"]: notify_cfg = ActionConfig({ "name": cfg["name"], @@ -366,32 +376,32 @@ class Quarantine: "pretend": cfg["pretend"], "type": "notify", "args": cfg["args"]["notify"].get_config()}) - self.notify = Notify(notify_cfg, local_addrs, debug) + self._notification = Notify(notify_cfg, local_addrs, debug) - self.whitelist = None + self._whitelist = None if "whitelist" in cfg["args"]: whitelist_cfg = cfg["args"]["whitelist"] whitelist_cfg["name"] = cfg["name"] whitelist_cfg["loglevel"] = cfg["loglevel"] - self.whitelist = Conditions( + self._whitelist = Conditions( whitelist_cfg, local_addrs=[], debug=debug) - self.milter_action = None + self._milter_action = None if "milter_action" in cfg["args"]: - self.milter_action = cfg["args"]["milter_action"] - self.reject_reason = None + self._milter_action = cfg["args"]["milter_action"] + self._reason = None if "reject_reason" in cfg["args"]: - self.reject_reason = cfg["args"]["reject_reason"] + self._reason = cfg["args"]["reject_reason"] def __str__(self): cfg = [] - cfg.append(f"store={str(self.store)}") - if self.notify is not None: - cfg.append(f"notify={str(self.notify)}") - if self.whitelist is not None: - cfg.append(f"whitelist={str(self.whitelist)}") + cfg.append(f"store={str(self._storage)}") + if self._notification is not None: + cfg.append(f"notify={str(self._notification)}") + if self._whitelist is not None: + cfg.append(f"whitelist={str(self._whitelist)}") for key in ["milter_action", "reject_reason"]: if key not in self.cfg["args"]: continue @@ -400,12 +410,79 @@ class Quarantine: class_name = type(self).__name__ return f"{class_name}(" + ", ".join(cfg) + ")" + @property + def name(self): + return self.cfg["name"] + + @property + def storage(self): + return self._storage.get_storage() + + @property + def notification(self): + if self._notification is None: + return None + return self._notification.get_notification() + + @property + def whitelist(self): + if self._whitelist is None: + return None + return self._whitelist.get_whitelist() + + @property + def milter_action(self): + return self._milter_action + + def notify(self, storage_id, recipient=None): + "Notify recipient about email in storage." + if not self._notification: + raise RuntimeError( + "notification not defined, " + "unable to send notification") + metadata, msg = self.storage.get_mail(storage_id) + + if recipient is not None: + if recipient not in metadata["recipients"]: + raise RuntimeError(f"invalid recipient '{recipient}'") + recipients = [recipient] + else: + recipients = metadata["recipients"] + + self.notification.notify(msg, metadata["queue_id"], + metadata["mailfrom"], recipients, + self.logger, metadata["vars"], + synchronous=True) + + def release(self, storage_id, recipients=None): + metadata, msg = self.storage.get_mail(storage_id) + if recipients and type(recipients) == str: + recipients = [recipients] + else: + recipients = metadata["recipients"] + + for recipient in recipients: + if recipient not in metadata["recipients"]: + raise RuntimeError(f"invalid recipient '{recipient}'") + try: + mailer.smtp_send( + self.smtp_host, + self.smtp_port, + metadata["mailfrom"], + recipient, + msg.as_string()) + + except Exception as e: + raise RuntimeError( + f"error while sending email to '{recipient}': {e}") + self.storage.delete(storage_id, recipient) + def execute(self, milter): logger = CustomLogger( self.logger, {"name": self.cfg["name"], "qid": milter.qid}) wl_rcpts = [] - if self.whitelist: - wl_rcpts = self.whitelist.get_wl_rcpts( + if self._whitelist: + wl_rcpts = self._whitelist.get_wl_rcpts( milter.msginfo["mailfrom"], milter.msginfo["rcpts"], logger) logger.info(f"whitelisted recipients: {wl_rcpts}") @@ -419,13 +496,13 @@ class Quarantine: logger.info(f"add to quarantine for recipients: {rcpts}") milter.msginfo["rcpts"] = rcpts - self.store.execute(milter) + self._storage.execute(milter) - if self.notify is not None: - self.notify.execute(milter) + if self._notification is not None: + 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"]: - return (self.milter_action, self.reject_reason) + if self._milter_action is not None and not milter.msginfo["rcpts"]: + return (self._milter_action, self._reason)