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

fix update logic from msg headers to milter headers

Browse Source
This commit is contained in:
Thomas Oettli
2020-11-13 17:32:45 +01:00
parent 267d0dbf1f
commit 6dbc50385f
1 changed files with 73 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

92
pymodmilter/actions.py
View File

@@ -35,10 +35,16 @@ def _replace_illegal_chars(string):
"\n", "") "\n", "")
def add_header(milter, field, value, pretend=False, update_msg=True, def _add_header(milter, field, value, idx=-1):
value = _replace_illegal_chars(
Header(s=value).encode())
milter.logger.debug(f"milter: addheader: {field}: {value}")
milter.addheader(field, value, idx)
def add_header(milter, field, value, pretend=False,
logger=logging.getLogger(__name__)): logger=logging.getLogger(__name__)):
"""Add a mail header field.""" """Add a mail header field."""
if update_msg:
header = f"{field}: {value}" header = f"{field}: {value}"
if logger.getEffectiveLevel() == logging.DEBUG: if logger.getEffectiveLevel() == logging.DEBUG:
logger.debug(f"add_header: {header}") logger.debug(f"add_header: {header}")
@@ -50,14 +56,19 @@ def add_header(milter, field, value, pretend=False, update_msg=True,
if pretend: if pretend:
return return
encoded_value = _replace_illegal_chars( _add_header(milter, field, value)
def _mod_header(milter, field, value, occ=1):
value = _replace_illegal_chars(
Header(s=value).encode()) Header(s=value).encode())
milter.logger.debug(f"milter: addheader: {field}: {encoded_value}") milter.logger.debug(
milter.addheader(field, encoded_value, -1) f"milter: chgheader: {field}[{occ}]: {value}")
milter.chgheader(field, occ, value)
def mod_header(milter, field, value, search=None, pretend=False, def mod_header(milter, field, value, search=None, pretend=False,
update_msg=True, logger=logging.getLogger(__name__)): logger=logging.getLogger(__name__)):
"""Change the value of a mail header field.""" """Change the value of a mail header field."""
if isinstance(field, str): if isinstance(field, str):
field = re.compile(field, re.IGNORECASE) field = re.compile(field, re.IGNORECASE)
@@ -74,22 +85,23 @@ def mod_header(milter, field, value, search=None, pretend=False,
if not field.match(f): if not field.match(f):
continue continue
new_v = v new_value = v
if search is not None: if search is not None:
new_v = search.sub(value, v).strip() new_value = search.sub(value, v).strip()
else:
new_value = value
if new_v == v: if not new_value:
continue
if not new_v:
logger.warning( logger.warning(
"mod_header: resulting value is empty, " "mod_header: resulting value is empty, "
"skip modification") "skip modification")
continue continue
if update_msg: if new_value == v:
continue
header = f"{f}: {v}" header = f"{f}: {v}"
new_header = f"{f}: {new_v}" new_header = f"{f}: {new_value}"
if logger.getEffectiveLevel() == logging.DEBUG: if logger.getEffectiveLevel() == logging.DEBUG:
logger.debug(f"mod_header: {header}: {new_header}") logger.debug(f"mod_header: {header}: {new_header}")
@@ -97,19 +109,21 @@ def mod_header(milter, field, value, search=None, pretend=False,
logger.info(f"mod_header: {header[0:70]}: {new_header[0:70]}") logger.info(f"mod_header: {header[0:70]}: {new_header[0:70]}")
milter.msg.replace_header( milter.msg.replace_header(
f, _replace_illegal_chars(new_v), occ=occ[f_lower]) f, _replace_illegal_chars(new_value), occ=occ[f_lower])
if pretend: if pretend:
continue continue
encoded_value = _replace_illegal_chars( _mod_header(milter, f, new_value, occ=occ[f_lower])
Header(s=new_v).encode())
def _del_header(milter, field, occ=1):
milter.logger.debug( milter.logger.debug(
f"milter: chgheader: {f}[{occ[f_lower]}]: {encoded_value}") f"milter: delheader: {field}[{occ}]")
milter.chgheader(f, occ[f_lower], encoded_value) milter.chgheader(field, occ, "")
def del_header(milter, field, value=None, pretend=False, update_msg=True, def del_header(milter, field, value=None, pretend=False,
logger=logging.getLogger(__name__)): logger=logging.getLogger(__name__)):
"""Delete a mail header field.""" """Delete a mail header field."""
if isinstance(field, str): if isinstance(field, str):
@@ -130,7 +144,6 @@ def del_header(milter, field, value=None, pretend=False, update_msg=True,
if value is not None and not value.search(v): if value is not None and not value.search(v):
continue continue
if update_msg:
header = f"{f}: {v}" header = f"{f}: {v}"
if logger.getEffectiveLevel() == logging.DEBUG: if logger.getEffectiveLevel() == logging.DEBUG:
logger.debug(f"del_header: {header}") logger.debug(f"del_header: {header}")
@@ -138,16 +151,14 @@ def del_header(milter, field, value=None, pretend=False, update_msg=True,
logger.info(f"del_header: {header[0:70]}") logger.info(f"del_header: {header[0:70]}")
milter.msg.remove_header(f, occ=occ[f_lower]) milter.msg.remove_header(f, occ=occ[f_lower])
occ[f_lower] -= 1
if not pretend: if not pretend:
milter.logger.debug( _del_header(milter, f, occ=occ[f_lower])
f"milter: chgheader: {f}[{occ[f_lower]}]:")
milter.chgheader(f, occ[f_lower], "") occ[f_lower] -= 1
def _serialize_msg(msg, logger): def _serialize_msg(msg, logger):
if not msg["MIME-Version"]: if msg.is_multipart() and not msg["MIME-Version"]:
msg.add_header("MIME-Version", "1.0") msg.add_header("MIME-Version", "1.0")
try: try:
@@ -242,21 +253,26 @@ def _update_body(milter, logger):
def _update_headers(milter, original_headers, logger): def _update_headers(milter, original_headers, logger):
# serialize the message object so it updates its headers internally if milter.msg.is_multipart() and not milter.msg["MIME-Version"]:
milter.msg.add_header("MIME-Version", "1.0")
# serialize the message object so it updates its internal strucure
milter.msg.as_bytes() milter.msg.as_bytes()
for field, value in original_headers:
if field not in milter.msg: original_headers = [(f, f.lower(), v) for f, v in original_headers]
del_header(milter, field=f"^{field}$", update_msg=False, headers = [(f, f.lower(), v) for f, v in milter.msg.items()]
logger=logger)
occ = defaultdict(int)
for field, field_lower, value in original_headers:
occ[field_lower] += 1
if (field, field_lower, value) not in headers:
_del_header(milter, field, occ=occ[field_lower])
occ[field] -= 1
for field, value in milter.msg.items(): for field, value in milter.msg.items():
field_lower = field.lower() field_lower = field.lower()
if not [f for f in original_headers if f[0].lower() == field_lower]: if (field, field_lower, value) not in original_headers:
add_header(milter, field=field, value=value, update_msg=False, _add_header(milter, field, value)
logger=logger)
else:
mod_header(milter, field=f"^{field}$", value=value,
update_msg=False, logger=logger)
def _wrap_message(milter, logger): def _wrap_message(milter, logger):