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

fix replacebody logic

Browse Source
This commit is contained in:
Thomas Oettli
2021-09-21 01:35:47 +02:00
parent 6333a0913d
commit f4bc545f9b
1 changed files with 18 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
pymodmilter/__init__.py
View File

@@ -190,7 +190,16 @@ class ModifyMilter(Milter.Base):
self.addheader(field, value) self.addheader(field, value)
def replacebody(self): def replacebody(self):
self._replacebody = True self._body_changed = True
def _replacebody(self):
if not self._body_changed:
return
data = self.msg.as_bytes()
body_pos = data.find(b"\r\n\r\n") + 4
self.logger.debug("replace body")
super().replacebody(data[body_pos:])
del data
def delrcpt(self, rcpts): def delrcpt(self, rcpts):
"Remove recipient. May be called from eom callback only." "Remove recipient. May be called from eom callback only."
@@ -354,7 +363,7 @@ class ModifyMilter(Milter.Base):
"rcpts": [*self.rcpts], "rcpts": [*self.rcpts],
"vars": {}} "vars": {}}
self._replacebody = False self._body_changed = False
milter_action = None milter_action = None
for rule in self.rules: for rule in self.rules:
milter_action = rule.execute(self) milter_action = rule.execute(self)
@@ -363,23 +372,16 @@ class ModifyMilter(Milter.Base):
if milter_action is not None: if milter_action is not None:
break break
if self._replacebody: if milter_action is None:
data = self.msg.as_bytes() self._replacebody()
body_pos = data.find(b"\r\n\r\n") + 4 else:
self.logger.debug("replace body")
super().replacebody(data[body_pos:])
del data
if milter_action is not None:
action, reason = milter_action action, reason = milter_action
if action == "REJECT": if action == "ACCEPT":
self._replacebody()
elif action == "REJECT":
self.setreply("554", "5.7.0", reason) self.setreply("554", "5.7.0", reason)
return Milter.REJECT return Milter.REJECT
elif action == "DISCARD":
if action == "ACCEPT":
return Milter.ACCEPT
if action == "DISCARD":
return Milter.DISCARD return Milter.DISCARD
except Exception as e: except Exception as e: