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

fix bugs

Browse Source
This commit is contained in:
Thomas Oettli
2023-12-12 18:15:09 +01:00
parent af800c73aa
commit a226bc70a9
4 changed files with 109 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
pyquarantine/__init__.py
View File

@@ -71,7 +71,7 @@ class QuarantineMilter(Milter.Base):
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
logger.setLevel(QuarantineMilter._loglevel) 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) rule = Rule(rule_cfg, local_addrs, debug)
logger.debug(rule) logger.debug(rule)
QuarantineMilter._rules.append(rule) QuarantineMilter._rules.append(rule)

85
pyquarantine/cli.py
View File

@@ -21,29 +21,41 @@ import logging.handlers
import sys import sys
import time 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.storage import Quarantine
from pyquarantine import __version__ as version 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: 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: except StopIteration:
raise RuntimeError(f"invalid quarantine '{name}'") 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): def _get_notification(cfg, name, debug):
notification = _get_quarantine(quarantines, name, debug).notification notification = _get_quarantine(cfg, name, debug).notification
if not notification: if not notification:
raise RuntimeError( raise RuntimeError(
"notification type is set to NONE") "notification type is set to NONE")
return notification return notification
def _get_allowlist(quarantines, name, debug): def _get_allowlist(cfg, name, debug):
allowlist = _get_quarantine(quarantines, name, debug).allowlist allowlist = _get_quarantine(cfg, name, debug).allowlist
if not allowlist: if not allowlist:
raise RuntimeError( raise RuntimeError(
"allowlist type is set to NONE") "allowlist type is set to NONE")
@@ -91,7 +103,8 @@ def print_table(columns, rows):
print(row_format.format(*row)) print(row_format.format(*row))
def list_quarantines(quarantines, args): def list_quarantines(cfg, args):
quarantines = _get_quarantines(cfg)
if args.batch: if args.batch:
print("\n".join([q["name"] for q in quarantines])) print("\n".join([q["name"] for q in quarantines]))
else: else:
@@ -132,8 +145,8 @@ def list_quarantines(quarantines, args):
) )
def list_quarantine_emails(quarantines, args): def list_quarantine_emails(cfg, args):
storage = _get_quarantine(quarantines, args.quarantine, args.debug).storage storage = _get_quarantine(cfg, args.quarantine, args.debug).storage
# find emails and transform some metadata values to strings # find emails and transform some metadata values to strings
rows = [] rows = []
@@ -180,8 +193,8 @@ def list_quarantine_emails(quarantines, args):
) )
def list_allowlist(quarantines, args): def list_allowlist(cfg, args):
allowlist = _get_allowlist(quarantines, args.quarantine, args.debug) allowlist = _get_allowlist(cfg, args.quarantine, args.debug)
# find allowlist entries # find allowlist entries
entries = allowlist.find( 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__) logger = logging.getLogger(__name__)
allowlist = _get_allowlist(quarantines, args.quarantine, args.debug) allowlist = _get_allowlist(cfg, args.quarantine, args.debug)
# check existing entries # check existing entries
entries = allowlist.check(args.mailfrom, args.recipient, logger) entries = allowlist.check(args.mailfrom, args.recipient, logger)
@@ -250,21 +263,21 @@ def add_allowlist_entry(quarantines, args):
print("allowlist entry added successfully") print("allowlist entry added successfully")
def delete_allowlist_entry(quarantines, args): def delete_allowlist_entry(cfg, args):
allowlist = _get_allowlist(quarantines, args.quarantine, args.debug) allowlist = _get_allowlist(cfg, args.quarantine, args.debug)
allowlist.delete(args.allowlist_id) allowlist.delete(args.allowlist_id)
print("allowlist entry deleted successfully") print("allowlist entry deleted successfully")
def notify(quarantines, args): def notify(cfg, args):
quarantine = _get_quarantine(quarantines, args.quarantine, args.debug) quarantine = _get_quarantine(cfg, args.quarantine, args.debug)
quarantine.notify(args.quarantine_id, args.recipient) quarantine.notify(args.quarantine_id, args.recipient)
print("notification sent successfully") print("notification sent successfully")
def release(quarantines, args): def release(cfg, args):
logger = logging.getLogger(__name__) 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 = quarantine.release(args.quarantine_id, args.recipient)
rcpts = ", ".join(rcpts) rcpts = ", ".join(rcpts)
logger.info( logger.info(
@@ -272,29 +285,29 @@ def release(quarantines, args):
f"for {rcpts}") f"for {rcpts}")
def copy(quarantines, args): def copy(cfg, args):
logger = logging.getLogger(__name__) 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) quarantine.copy(args.quarantine_id, args.recipient)
logger.info( logger.info(
f"{args.quarantine}: sent a copy of message with id {args.quarantine_id} " f"{args.quarantine}: sent a copy of message with id "
f"to {args.recipient}") f"{args.quarantine_id} to {args.recipient}")
def delete(quarantines, args): def delete(cfg, args):
storage = _get_quarantine(quarantines, args.quarantine, args.debug).storage storage = _get_quarantine(cfg, args.quarantine, args.debug).storage
storage.delete(args.quarantine_id, args.recipient) storage.delete(args.quarantine_id, args.recipient)
print("quarantined message deleted successfully") print("quarantined message deleted successfully")
def get(quarantines, args): def get(cfg, args):
storage = _get_quarantine(quarantines, args.quarantine, args.debug).storage storage = _get_quarantine(cfg, args.quarantine, args.debug).storage
data = storage.get_mail_bytes(args.quarantine_id) data = storage.get_mail_bytes(args.quarantine_id)
sys.stdout.buffer.write(data) sys.stdout.buffer.write(data)
def metadata(quarantines, args): def metadata(cfg, args):
storage = _get_quarantine(quarantines, args.quarantine, args.debug).storage storage = _get_quarantine(cfg, args.quarantine, args.debug).storage
metadata = storage.get_metadata(args.quarantine_id) metadata = storage.get_metadata(args.quarantine_id)
print(json.dumps(metadata)) print(json.dumps(metadata))
@@ -622,6 +635,7 @@ def main():
try: try:
logger.debug("read milter configuration") logger.debug("read milter configuration")
cfg = get_milter_config(args.config, raw=True) cfg = get_milter_config(args.config, raw=True)
if "rules" not in cfg or not cfg["rules"]: if "rules" not in cfg or not cfg["rules"]:
raise RuntimeError("no rules configured") raise RuntimeError("no rules configured")
@@ -629,16 +643,11 @@ def main():
if "actions" not in rule or not rule["actions"]: if "actions" not in rule or not rule["actions"]:
raise RuntimeError( raise RuntimeError(
f"{rule['name']}: no actions configured") f"{rule['name']}: no actions configured")
except (RuntimeError, AssertionError) as e: except (RuntimeError, AssertionError) as e:
logger.error(f"config error: {e}") logger.error(f"config error: {e}")
sys.exit(255) sys.exit(255)
quarantines = []
for rule in cfg["rules"]:
for action in rule["actions"]:
if action["type"] == "quarantine":
quarantines.append(action)
if args.syslog: if args.syslog:
# setup syslog # setup syslog
sysloghandler = logging.handlers.SysLogHandler( sysloghandler = logging.handlers.SysLogHandler(
@@ -655,7 +664,7 @@ def main():
# call the commands function # call the commands function
try: try:
args.func(quarantines, args) args.func(cfg, args)
except RuntimeError as e: except RuntimeError as e:
logger.error(e) logger.error(e)
sys.exit(1) sys.exit(1)

74
pyquarantine/config.py
View File

@@ -43,7 +43,7 @@ class BaseConfig:
"properties": { "properties": {
"loglevel": {"type": "string", "default": "info"}}} "loglevel": {"type": "string", "default": "info"}}}
def __init__(self, config): def __init__(self, config, lists):
required = self.JSON_SCHEMA["required"] required = self.JSON_SCHEMA["required"]
properties = self.JSON_SCHEMA["properties"] properties = self.JSON_SCHEMA["properties"]
for p in properties.keys(): for p in properties.keys():
@@ -92,16 +92,18 @@ class BaseConfig:
class ListConfig(BaseConfig): class ListConfig(BaseConfig):
JSON_SCHEMA = { JSON_SCHEMA = {
"type": "object", "type": "object",
"required": ["type"], "required": ["name", "type"],
"additionalProperties": True, "additionalProperties": True,
"properties": { "properties": {
"type": {"enum": ["db"]}}, "type": {"enum": ["db"]},
"name": {"type": "string"}},
"if": {"properties": {"type": {"const": "db"}}}, "if": {"properties": {"type": {"const": "db"}}},
"then": { "then": {
"required": ["connection", "table"], "required": ["connection", "table"],
"additionalProperties": False, "additionalProperties": False,
"properties": { "properties": {
"type": {"type": "string"}, "type": {"type": "string"},
"name": {"type": "string"},
"connection": {"type": "string"}, "connection": {"type": "string"},
"table": {"type": "string"}}}} "table": {"type": "string"}}}}
@@ -121,14 +123,18 @@ class ConditionsConfig(BaseConfig):
"headers": {"type": "array", "headers": {"type": "array",
"items": {"type": "string"}}, "items": {"type": "string"}},
"var": {"type": "string"}, "var": {"type": "string"},
"list": {"type": "object"}}} "list": {"type": "string"}}}
def __init__(self, config, rec=True): def __init__(self, config, lists, rec=True):
super().__init__(config) 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: if not rec:
return return
if "list" in self:
self["list"] = ListConfig(self["list"])
class AddHeaderConfig(BaseConfig): class AddHeaderConfig(BaseConfig):
@@ -242,22 +248,26 @@ class QuarantineConfig(BaseConfig):
"notify": {"type": "object"}, "notify": {"type": "object"},
"milter_action": {"type": "string"}, "milter_action": {"type": "string"},
"reject_reason": {"type": "string"}, "reject_reason": {"type": "string"},
"allowlist": {"type": "object"}, "allowlist": {"type": "string"},
"store": {"type": "object"}, "store": {"type": "object"},
"smtp_host": {"type": "string"}, "smtp_host": {"type": "string"},
"smtp_port": {"type": "number"}}} "smtp_port": {"type": "number"}}}
def __init__(self, config, rec=True): def __init__(self, config, lists, rec=True):
super().__init__(config) super().__init__(config, lists)
if not rec: if not rec:
return return
if "metadata" not in self["store"]: if "metadata" not in self["store"]:
self["store"]["metadata"] = True self["store"]["metadata"] = True
self["store"] = StoreConfig(self["store"]) self["store"] = StoreConfig(self["store"], lists)
if "notify" in self: if "notify" in self:
self["notify"] = NotifyConfig(self["notify"]) self["notify"] = NotifyConfig(self["notify"], lists)
if "allowlist" in self: 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): class ActionConfig(BaseConfig):
@@ -283,13 +293,13 @@ class ActionConfig(BaseConfig):
"type": {"enum": list(ACTION_TYPES.keys())}, "type": {"enum": list(ACTION_TYPES.keys())},
"options": {"type": "object"}}} "options": {"type": "object"}}}
def __init__(self, config, rec=True): def __init__(self, config, lists, rec=True):
super().__init__(config) super().__init__(config, lists)
if not rec: if not rec:
return return
if "conditions" in self: if "conditions" in self:
self["conditions"] = ConditionsConfig(self["conditions"]) self["conditions"] = ConditionsConfig(self["conditions"], lists)
self["action"] = self.ACTION_TYPES[self["type"]](self["options"]) self["action"] = self.ACTION_TYPES[self["type"]](self["options"], lists)
class RuleConfig(BaseConfig): class RuleConfig(BaseConfig):
@@ -304,20 +314,20 @@ class RuleConfig(BaseConfig):
"conditions": {"type": "object"}, "conditions": {"type": "object"},
"actions": {"type": "array"}}} "actions": {"type": "array"}}}
def __init__(self, config, rec=True): def __init__(self, config, lists, rec=True):
super().__init__(config) super().__init__(config, lists)
if not rec: if not rec:
return return
if "conditions" in self: if "conditions" in self:
self["conditions"] = ConditionsConfig(self["conditions"]) self["conditions"] = ConditionsConfig(self["conditions"], lists)
actions = [] actions = []
for idx, action in enumerate(self["actions"]): for action in self["actions"]:
if "loglevel" not in action: if "loglevel" not in action:
action["loglevel"] = config["loglevel"] action["loglevel"] = config["loglevel"]
if "pretend" not in action: if "pretend" not in action:
action["pretend"] = config["pretend"] action["pretend"] = config["pretend"]
actions.append(ActionConfig(action, rec)) actions.append(ActionConfig(action, lists, rec))
self["actions"] = actions self["actions"] = actions
@@ -339,19 +349,25 @@ class QuarantineMilterConfig(BaseConfig):
"192.168.0.0/16"]}, "192.168.0.0/16"]},
"loglevel": {"type": "string", "default": "info"}, "loglevel": {"type": "string", "default": "info"},
"pretend": {"type": "boolean", "default": False}, "pretend": {"type": "boolean", "default": False},
"rules": {"type": "array"}}} "rules": {"type": "array"},
"lists": {"type": "array",
"default": []}}}
def __init__(self, config, rec=True): def __init__(self, config, rec=True):
super().__init__(config) super().__init__(config, {})
if not rec: if not rec:
return return
lists = {}
for lst in self["lists"]:
lists[lst["name"]] = ListConfig(lst, rec)
self["lists"] = lists
rules = [] rules = []
for idx, rule in enumerate(self["rules"]): for rule in self["rules"]:
if "loglevel" not in rule: if "loglevel" not in rule:
rule["loglevel"] = config["loglevel"] rule["loglevel"] = config["loglevel"]
if "pretend" not in rule: if "pretend" not in rule:
rule["pretend"] = config["pretend"] rule["pretend"] = config["pretend"]
rules.append(RuleConfig(rule, rec)) rules.append(RuleConfig(rule, lists, rec))
self["rules"] = rules self["rules"] = rules
@@ -371,5 +387,9 @@ def get_milter_config(cfgfile, raw=False):
msg = "\n".join(cfg_text) msg = "\n".join(cfg_text)
raise RuntimeError(f"{e}\n{msg}") raise RuntimeError(f"{e}\n{msg}")
if raw: if raw:
lists = {}
for lst in cfg["lists"]:
lists[lst["name"]] = lst
cfg["lists"] = lists
return cfg return cfg
return QuarantineMilterConfig(cfg) return QuarantineMilterConfig(cfg)

9
pyquarantine/storage.py
View File

@@ -407,12 +407,14 @@ class Quarantine:
self.logger = logging.getLogger(cfg["name"]) self.logger = logging.getLogger(cfg["name"])
self.logger.setLevel(cfg.get_loglevel(debug)) self.logger.setLevel(cfg.get_loglevel(debug))
storage_cfg = ActionConfig({ storage_cfg = ActionConfig(
{
"name": cfg["name"], "name": cfg["name"],
"loglevel": cfg["loglevel"], "loglevel": cfg["loglevel"],
"pretend": cfg["pretend"], "pretend": cfg["pretend"],
"type": "store", "type": "store",
"options": cfg["options"]["store"].get_config()}) "options": cfg["options"]["store"].get_config()},
{})
self._storage = Store(storage_cfg, local_addrs, debug) self._storage = Store(storage_cfg, local_addrs, debug)
self.smtp_host = cfg["options"]["smtp_host"] self.smtp_host = cfg["options"]["smtp_host"]
@@ -425,7 +427,8 @@ class Quarantine:
"loglevel": cfg["loglevel"], "loglevel": cfg["loglevel"],
"pretend": cfg["pretend"], "pretend": cfg["pretend"],
"type": "notify", "type": "notify",
"options": cfg["options"]["notify"].get_config()}) "options": cfg["options"]["notify"].get_config()},
{})
self._notification = Notify(notify_cfg, local_addrs, debug) self._notification = Notify(notify_cfg, local_addrs, debug)
self._allowlist = None self._allowlist = None