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

fix file loading and improve quarantine logic

Browse Source
This commit is contained in:
Thomas Oettli
2021-10-04 19:03:55 +02:00
parent 5c9ec5037e
commit 3220fcbf2c
2 changed files with 31 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
pyquarantine/notify.py
View File

@@ -121,20 +121,21 @@ class EMailNotification(BaseNotification):
self.from_header = from_header self.from_header = from_header
self.subject = subject self.subject = subject
try: try:
self.template = open(template, "r").read() with open(template, "r") as fh:
self.template = fh.read()
self.embed_imgs = [] self.embed_imgs = []
for img_path in embed_imgs: for img_path in embed_imgs:
img = MIMEImage(open(img_path, "rb").read()) with open(img_path, "rb") as fh:
img = MIMEImage(fh.read())
filename = basename(img_path) filename = basename(img_path)
img.add_header("Content-ID", f"<{filename}>") img.add_header("Content-ID", f"<{filename}>")
self.embed_imgs.append(img) self.embed_imgs.append(img)
self.replacement_img = repl_img self.replacement_img = repl_img
self.strip_images = strip_imgs self.strip_images = strip_imgs
if not strip_imgs and repl_img: if not strip_imgs and repl_img:
self.replacement_img = MIMEImage( with open(repl_img, "rb") as fh:
open(repl_img, "rb").read()) self.replacement_img = MIMEImage(fh.read())
self.replacement_img.add_header( self.replacement_img.add_header(
"Content-ID", "<removed_for_security_reasons>") "Content-ID", "<removed_for_security_reasons>")

26
pyquarantine/storage.py
View File

@@ -177,7 +177,8 @@ class FileMailStorage(BaseMailStorage):
# save mail # save mail
self._save_datafile(datafile, data) self._save_datafile(datafile, data)
logger.info(f"stored message in file {datafile}") logger.debug(f"stored message in file {datafile}")
logger.info(f"storaged message with id {storage_id}")
if not self.metadata: if not self.metadata:
return storage_id, None, datafile return storage_id, None, datafile
@@ -393,10 +394,16 @@ class Quarantine:
self._milter_action = None self._milter_action = None
if "milter_action" in cfg["args"]: if "milter_action" in cfg["args"]:
self._milter_action = cfg["args"]["milter_action"] self._milter_action = cfg["args"]["milter_action"].upper()
assert self._milter_action in ["ACCEPT", "REJECT"], \
f"invalid milter_action '{cfg['args']['milter_action']}'"
self._reason = None self._reason = None
if self._milter_action == "REJECT":
if "reject_reason" in cfg["args"]: if "reject_reason" in cfg["args"]:
self._reason = cfg["args"]["reject_reason"] self._reason = cfg["args"]["reject_reason"]
else:
self._reason = "Message rejected"
def __str__(self): def __str__(self):
cfg = [] cfg = []
@@ -483,19 +490,17 @@ class Quarantine:
def execute(self, milter): def execute(self, milter):
logger = CustomLogger( logger = CustomLogger(
self.logger, {"name": self.cfg["name"], "qid": milter.qid}) self.logger, {"name": self.cfg["name"], "qid": milter.qid})
rcpts = milter.msginfo["rcpts"]
wl_rcpts = [] wl_rcpts = []
if self._whitelist: if self._whitelist:
wl_rcpts = self._whitelist.get_wl_rcpts( wl_rcpts = self._whitelist.get_wl_rcpts(
milter.msginfo["mailfrom"], milter.msginfo["rcpts"], logger) milter.msginfo["mailfrom"], rcpts, logger)
if wl_rcpts:
logger.info(f"whitelisted recipients: {wl_rcpts}") logger.info(f"whitelisted recipients: {wl_rcpts}")
rcpts = [rcpt for rcpt in rcpts if rcpt not in wl_rcpts]
rcpts = [
rcpt for rcpt in milter.msginfo["rcpts"] if rcpt not in wl_rcpts]
if not rcpts: if not rcpts:
# all recipients whitelisted # all recipients whitelisted
return return
logger.info(f"add to quarantine for recipients: {rcpts}") logger.info(f"add to quarantine for recipients: {rcpts}")
milter.msginfo["rcpts"] = rcpts milter.msginfo["rcpts"] = rcpts
@@ -505,7 +510,8 @@ class Quarantine:
self._notification.execute(milter) self._notification.execute(milter)
milter.msginfo["rcpts"].extend(wl_rcpts) milter.msginfo["rcpts"].extend(wl_rcpts)
milter.delrcpt(rcpts)
if self._milter_action is not None and not milter.msginfo["rcpts"]: if self._milter_action is not None:
milter.delrcpt(rcpts)
if not milter.msginfo["rcpts"]:
return (self._milter_action, self._reason) return (self._milter_action, self._reason)