adapt all CLI functions to new code structure

This commit is contained in:
2021-09-30 23:40:50 +02:00
parent 9e0baf3ce9
commit 01ae131088
7 changed files with 170 additions and 109 deletions

View File

@@ -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)

View File

@@ -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)

View File

@@ -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"]})

View File

@@ -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

View File

@@ -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})

View File

@@ -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):

View File

@@ -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)