From a226bc70a9f4a14cbdd0d6f0bfaab88e48b66ea2 Mon Sep 17 00:00:00 2001 From: Thomas Oettli Date: Tue, 12 Dec 2023 18:15:09 +0100 Subject: [PATCH] fix bugs --- pyquarantine/__init__.py | 2 +- pyquarantine/cli.py | 85 ++++++++++++++++++++++------------------ pyquarantine/config.py | 74 +++++++++++++++++++++------------- pyquarantine/storage.py | 25 ++++++------ 4 files changed, 109 insertions(+), 77 deletions(-) diff --git a/pyquarantine/__init__.py b/pyquarantine/__init__.py index 993bf57..380c8b1 100644 --- a/pyquarantine/__init__.py +++ b/pyquarantine/__init__.py @@ -71,7 +71,7 @@ class QuarantineMilter(Milter.Base): logger = logging.getLogger(__name__) logger.setLevel(QuarantineMilter._loglevel) - for idx, rule_cfg in enumerate(cfg["rules"]): + for rule_cfg in cfg["rules"]: rule = Rule(rule_cfg, local_addrs, debug) logger.debug(rule) QuarantineMilter._rules.append(rule) diff --git a/pyquarantine/cli.py b/pyquarantine/cli.py index 693d262..1fde908 100644 --- a/pyquarantine/cli.py +++ b/pyquarantine/cli.py @@ -21,29 +21,41 @@ import logging.handlers import sys import time -from pyquarantine.config import get_milter_config, ActionConfig +from pyquarantine.config import get_milter_config, ActionConfig, ListConfig from pyquarantine.storage import Quarantine from pyquarantine import __version__ as version -def _get_quarantine(quarantines, name, debug): +def _get_quarantines(cfg): + quarantines = [] + for rule in cfg["rules"]: + for action in rule["actions"]: + if action["type"] == "quarantine": + quarantines.append(action) + return quarantines + + +def _get_quarantine(cfg, name, debug): try: - quarantine = next((q for q in quarantines if q["name"] == name)) + quarantine = next( + (q for q in _get_quarantines(cfg) if q["name"] == name)) except StopIteration: raise RuntimeError(f"invalid quarantine '{name}'") - return Quarantine(ActionConfig(quarantine), [], debug) + for name, lst in cfg["lists"].items(): + cfg["lists"][name] = ListConfig(lst, {}) + return Quarantine(ActionConfig(quarantine, cfg["lists"]), [], debug) -def _get_notification(quarantines, name, debug): - notification = _get_quarantine(quarantines, name, debug).notification +def _get_notification(cfg, name, debug): + notification = _get_quarantine(cfg, name, debug).notification if not notification: raise RuntimeError( "notification type is set to NONE") return notification -def _get_allowlist(quarantines, name, debug): - allowlist = _get_quarantine(quarantines, name, debug).allowlist +def _get_allowlist(cfg, name, debug): + allowlist = _get_quarantine(cfg, name, debug).allowlist if not allowlist: raise RuntimeError( "allowlist type is set to NONE") @@ -91,7 +103,8 @@ def print_table(columns, rows): print(row_format.format(*row)) -def list_quarantines(quarantines, args): +def list_quarantines(cfg, args): + quarantines = _get_quarantines(cfg) if args.batch: print("\n".join([q["name"] for q in quarantines])) else: @@ -132,8 +145,8 @@ def list_quarantines(quarantines, args): ) -def list_quarantine_emails(quarantines, args): - storage = _get_quarantine(quarantines, args.quarantine, args.debug).storage +def list_quarantine_emails(cfg, args): + storage = _get_quarantine(cfg, args.quarantine, args.debug).storage # find emails and transform some metadata values to strings rows = [] @@ -180,8 +193,8 @@ def list_quarantine_emails(quarantines, args): ) -def list_allowlist(quarantines, args): - allowlist = _get_allowlist(quarantines, args.quarantine, args.debug) +def list_allowlist(cfg, args): + allowlist = _get_allowlist(cfg, args.quarantine, args.debug) # find allowlist entries entries = allowlist.find( @@ -210,9 +223,9 @@ def list_allowlist(quarantines, args): ) -def add_allowlist_entry(quarantines, args): +def add_allowlist_entry(cfg, args): logger = logging.getLogger(__name__) - allowlist = _get_allowlist(quarantines, args.quarantine, args.debug) + allowlist = _get_allowlist(cfg, args.quarantine, args.debug) # check existing entries entries = allowlist.check(args.mailfrom, args.recipient, logger) @@ -250,21 +263,21 @@ def add_allowlist_entry(quarantines, args): print("allowlist entry added successfully") -def delete_allowlist_entry(quarantines, args): - allowlist = _get_allowlist(quarantines, args.quarantine, args.debug) +def delete_allowlist_entry(cfg, args): + allowlist = _get_allowlist(cfg, args.quarantine, args.debug) allowlist.delete(args.allowlist_id) print("allowlist entry deleted successfully") -def notify(quarantines, args): - quarantine = _get_quarantine(quarantines, args.quarantine, args.debug) +def notify(cfg, args): + quarantine = _get_quarantine(cfg, args.quarantine, args.debug) quarantine.notify(args.quarantine_id, args.recipient) print("notification sent successfully") -def release(quarantines, args): +def release(cfg, args): logger = logging.getLogger(__name__) - quarantine = _get_quarantine(quarantines, args.quarantine, args.debug) + quarantine = _get_quarantine(cfg, args.quarantine, args.debug) rcpts = quarantine.release(args.quarantine_id, args.recipient) rcpts = ", ".join(rcpts) logger.info( @@ -272,29 +285,29 @@ def release(quarantines, args): f"for {rcpts}") -def copy(quarantines, args): +def copy(cfg, args): logger = logging.getLogger(__name__) - quarantine = _get_quarantine(quarantines, args.quarantine, args.debug) + quarantine = _get_quarantine(cfg, args.quarantine, args.debug) quarantine.copy(args.quarantine_id, args.recipient) logger.info( - f"{args.quarantine}: sent a copy of message with id {args.quarantine_id} " - f"to {args.recipient}") + f"{args.quarantine}: sent a copy of message with id " + f"{args.quarantine_id} to {args.recipient}") -def delete(quarantines, args): - storage = _get_quarantine(quarantines, args.quarantine, args.debug).storage +def delete(cfg, args): + storage = _get_quarantine(cfg, args.quarantine, args.debug).storage storage.delete(args.quarantine_id, args.recipient) print("quarantined message deleted successfully") -def get(quarantines, args): - storage = _get_quarantine(quarantines, args.quarantine, args.debug).storage +def get(cfg, args): + storage = _get_quarantine(cfg, args.quarantine, args.debug).storage data = storage.get_mail_bytes(args.quarantine_id) sys.stdout.buffer.write(data) -def metadata(quarantines, args): - storage = _get_quarantine(quarantines, args.quarantine, args.debug).storage +def metadata(cfg, args): + storage = _get_quarantine(cfg, args.quarantine, args.debug).storage metadata = storage.get_metadata(args.quarantine_id) print(json.dumps(metadata)) @@ -622,6 +635,7 @@ def main(): try: logger.debug("read milter configuration") cfg = get_milter_config(args.config, raw=True) + if "rules" not in cfg or not cfg["rules"]: raise RuntimeError("no rules configured") @@ -629,16 +643,11 @@ def main(): if "actions" not in rule or not rule["actions"]: raise RuntimeError( f"{rule['name']}: no actions configured") + except (RuntimeError, AssertionError) as e: logger.error(f"config error: {e}") sys.exit(255) - quarantines = [] - for rule in cfg["rules"]: - for action in rule["actions"]: - if action["type"] == "quarantine": - quarantines.append(action) - if args.syslog: # setup syslog sysloghandler = logging.handlers.SysLogHandler( @@ -655,7 +664,7 @@ def main(): # call the commands function try: - args.func(quarantines, args) + args.func(cfg, args) except RuntimeError as e: logger.error(e) sys.exit(1) diff --git a/pyquarantine/config.py b/pyquarantine/config.py index ed72bd6..3e6b571 100644 --- a/pyquarantine/config.py +++ b/pyquarantine/config.py @@ -43,7 +43,7 @@ class BaseConfig: "properties": { "loglevel": {"type": "string", "default": "info"}}} - def __init__(self, config): + def __init__(self, config, lists): required = self.JSON_SCHEMA["required"] properties = self.JSON_SCHEMA["properties"] for p in properties.keys(): @@ -92,16 +92,18 @@ class BaseConfig: class ListConfig(BaseConfig): JSON_SCHEMA = { "type": "object", - "required": ["type"], + "required": ["name", "type"], "additionalProperties": True, "properties": { - "type": {"enum": ["db"]}}, + "type": {"enum": ["db"]}, + "name": {"type": "string"}}, "if": {"properties": {"type": {"const": "db"}}}, "then": { "required": ["connection", "table"], "additionalProperties": False, "properties": { "type": {"type": "string"}, + "name": {"type": "string"}, "connection": {"type": "string"}, "table": {"type": "string"}}}} @@ -121,14 +123,18 @@ class ConditionsConfig(BaseConfig): "headers": {"type": "array", "items": {"type": "string"}}, "var": {"type": "string"}, - "list": {"type": "object"}}} + "list": {"type": "string"}}} - def __init__(self, config, rec=True): - super().__init__(config) + def __init__(self, config, lists, rec=True): + super().__init__(config, lists) + if "list" in self: + lst = self["list"] + try: + self["list"] = lists[lst] + except KeyError: + raise RuntimeError(f"list '{lst}' is not configured") if not rec: return - if "list" in self: - self["list"] = ListConfig(self["list"]) class AddHeaderConfig(BaseConfig): @@ -242,22 +248,26 @@ class QuarantineConfig(BaseConfig): "notify": {"type": "object"}, "milter_action": {"type": "string"}, "reject_reason": {"type": "string"}, - "allowlist": {"type": "object"}, + "allowlist": {"type": "string"}, "store": {"type": "object"}, "smtp_host": {"type": "string"}, "smtp_port": {"type": "number"}}} - def __init__(self, config, rec=True): - super().__init__(config) + def __init__(self, config, lists, rec=True): + super().__init__(config, lists) if not rec: return if "metadata" not in self["store"]: self["store"]["metadata"] = True - self["store"] = StoreConfig(self["store"]) + self["store"] = StoreConfig(self["store"], lists) if "notify" in self: - self["notify"] = NotifyConfig(self["notify"]) + self["notify"] = NotifyConfig(self["notify"], lists) if "allowlist" in self: - self["allowlist"] = ListConfig(self["allowlist"]) + allowlist = self["allowlist"] + try: + self["allowlist"] = lists[allowlist] + except KeyError: + raise RuntimeError(f"list '{allowlist}' is not configured") class ActionConfig(BaseConfig): @@ -283,13 +293,13 @@ class ActionConfig(BaseConfig): "type": {"enum": list(ACTION_TYPES.keys())}, "options": {"type": "object"}}} - def __init__(self, config, rec=True): - super().__init__(config) + def __init__(self, config, lists, rec=True): + super().__init__(config, lists) if not rec: return if "conditions" in self: - self["conditions"] = ConditionsConfig(self["conditions"]) - self["action"] = self.ACTION_TYPES[self["type"]](self["options"]) + self["conditions"] = ConditionsConfig(self["conditions"], lists) + self["action"] = self.ACTION_TYPES[self["type"]](self["options"], lists) class RuleConfig(BaseConfig): @@ -304,20 +314,20 @@ class RuleConfig(BaseConfig): "conditions": {"type": "object"}, "actions": {"type": "array"}}} - def __init__(self, config, rec=True): - super().__init__(config) + def __init__(self, config, lists, rec=True): + super().__init__(config, lists) if not rec: return if "conditions" in self: - self["conditions"] = ConditionsConfig(self["conditions"]) + self["conditions"] = ConditionsConfig(self["conditions"], lists) actions = [] - for idx, action in enumerate(self["actions"]): + for action in 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)) + actions.append(ActionConfig(action, lists, rec)) self["actions"] = actions @@ -339,19 +349,25 @@ class QuarantineMilterConfig(BaseConfig): "192.168.0.0/16"]}, "loglevel": {"type": "string", "default": "info"}, "pretend": {"type": "boolean", "default": False}, - "rules": {"type": "array"}}} + "rules": {"type": "array"}, + "lists": {"type": "array", + "default": []}}} def __init__(self, config, rec=True): - super().__init__(config) + super().__init__(config, {}) if not rec: return + lists = {} + for lst in self["lists"]: + lists[lst["name"]] = ListConfig(lst, rec) + self["lists"] = lists rules = [] - for idx, rule in enumerate(self["rules"]): + for rule in 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)) + rules.append(RuleConfig(rule, lists, rec)) self["rules"] = rules @@ -371,5 +387,9 @@ def get_milter_config(cfgfile, raw=False): msg = "\n".join(cfg_text) raise RuntimeError(f"{e}\n{msg}") if raw: + lists = {} + for lst in cfg["lists"]: + lists[lst["name"]] = lst + cfg["lists"] = lists return cfg return QuarantineMilterConfig(cfg) diff --git a/pyquarantine/storage.py b/pyquarantine/storage.py index cdc8f4f..0185d3b 100644 --- a/pyquarantine/storage.py +++ b/pyquarantine/storage.py @@ -407,12 +407,14 @@ class Quarantine: self.logger = logging.getLogger(cfg["name"]) self.logger.setLevel(cfg.get_loglevel(debug)) - storage_cfg = ActionConfig({ - "name": cfg["name"], - "loglevel": cfg["loglevel"], - "pretend": cfg["pretend"], - "type": "store", - "options": cfg["options"]["store"].get_config()}) + storage_cfg = ActionConfig( + { + "name": cfg["name"], + "loglevel": cfg["loglevel"], + "pretend": cfg["pretend"], + "type": "store", + "options": cfg["options"]["store"].get_config()}, + {}) self._storage = Store(storage_cfg, local_addrs, debug) self.smtp_host = cfg["options"]["smtp_host"] @@ -421,11 +423,12 @@ class Quarantine: self._notification = None if "notify" in cfg["options"]: notify_cfg = ActionConfig({ - "name": cfg["name"], - "loglevel": cfg["loglevel"], - "pretend": cfg["pretend"], - "type": "notify", - "options": cfg["options"]["notify"].get_config()}) + "name": cfg["name"], + "loglevel": cfg["loglevel"], + "pretend": cfg["pretend"], + "type": "notify", + "options": cfg["options"]["notify"].get_config()}, + {}) self._notification = Notify(notify_cfg, local_addrs, debug) self._allowlist = None