change everything to pythons new email lib

This commit is contained in:
2020-11-11 00:49:49 +01:00
parent 59486a2e18
commit d07bb965b3
2 changed files with 129 additions and 146 deletions

View File

@@ -26,35 +26,13 @@ import Milter
import logging import logging
from Milter.utils import parse_addr from Milter.utils import parse_addr
from email.charset import Charset from email.message import Message
from email.header import Header, decode_header from email.parser import BytesFeedParser
from io import BytesIO from email.policy import default as default_policy
from pymodmilter.conditions import Conditions from pymodmilter.conditions import Conditions
def make_header(decoded_seq, maxlinelen=None, header_name=None,
continuation_ws=' ', errors='strict'):
"""Create a Header from a sequence of pairs as returned by decode_header()
decode_header() takes a header value string and returns a sequence of
pairs of the format (decoded_string, charset) where charset is the string
name of the character set.
This function takes one of those sequence of pairs and returns a Header
instance. Optional maxlinelen, header_name, and continuation_ws are as in
the Header constructor.
"""
h = Header(maxlinelen=maxlinelen, header_name=header_name,
continuation_ws=continuation_ws)
for s, charset in decoded_seq:
# None means us-ascii but we can simply pass it on to h.append()
if charset is not None and not isinstance(charset, Charset):
charset = Charset(charset)
h.append(s, charset, errors=errors)
return h
class CustomLogger(logging.LoggerAdapter): class CustomLogger(logging.LoggerAdapter):
def process(self, msg, kwargs): def process(self, msg, kwargs):
if "name" in self.extra: if "name" in self.extra:
@@ -91,17 +69,15 @@ class Rule:
self.actions = actions self.actions = actions
self.pretend = pretend self.pretend = pretend
self._needs = [] self._need_body = False
for action in actions: for action in actions:
for need in action.needs(): if action.need_body():
if need not in self._needs: self._need_body = True
self._needs.append(need) break
self.logger.debug("needs: {}".format(", ".join(self._needs))) def need_body(self):
"""Return the if this rule needs the message body."""
def needs(self): return self._need_body
"""Return the needs of this rule."""
return self._needs
def ignores(self, host=None, envfrom=None, envto=None): def ignores(self, host=None, envfrom=None, envto=None):
args = {} args = {}
@@ -122,17 +98,53 @@ class Rule:
return True return True
def execute(self, milter, pretend=None): def execute(self, milter, msg, pretend=None):
"""Execute all actions of this rule.""" """Execute all actions of this rule."""
if pretend is None: if pretend is None:
pretend = self.pretend pretend = self.pretend
for action in self.actions: for action in self.actions:
milter_action = action.execute(milter) milter_action = action.execute(milter, msg, pretend=pretend)
if milter_action is not None: if milter_action is not None:
return milter_action return milter_action
class MilterMessage(Message):
def replace_header(self, _name, _value, occ=None):
_name = _name.lower()
counter = 0
for i, (k, v) in zip(range(len(self._headers)), self._headers):
if k.lower() == _name:
counter += 1
if not occ or counter == occ:
self._headers[i] = self.policy.header_store_parse(
k, _value)
break
else:
raise KeyError(_name)
def remove_header(self, name, occ=None):
name = name.lower()
newheaders = []
counter = 0
for k, v in self._headers:
if k.lower() == name:
counter += 1
if counter != occ:
newheaders.append((k, v))
else:
newheaders.append((k, v))
self._headers = newheaders
def remove_surrogates(string):
return string.encode(
"ascii", errors="surrogateescape").decode(
"ascii", errors="replace")
class ModifyMilter(Milter.Base): class ModifyMilter(Milter.Base):
"""ModifyMilter based on Milter.Base to implement milter communication""" """ModifyMilter based on Milter.Base to implement milter communication"""
@@ -175,7 +187,7 @@ class ModifyMilter(Milter.Base):
return Milter.ACCEPT return Milter.ACCEPT
except Exception as e: except Exception as e:
self.logger.exception( self.logger.exception(
f"an exception occured in connect function: {e}") f"an exception occured in connect method: {e}")
return Milter.TEMPFAIL return Milter.TEMPFAIL
return Milter.CONTINUE return Milter.CONTINUE
@@ -196,7 +208,7 @@ class ModifyMilter(Milter.Base):
self.recipients = set() self.recipients = set()
except Exception as e: except Exception as e:
self.logger.exception( self.logger.exception(
f"an exception occured in envfrom function: {e}") f"an exception occured in envfrom method: {e}")
return Milter.TEMPFAIL return Milter.TEMPFAIL
return Milter.CONTINUE return Milter.CONTINUE
@@ -207,7 +219,7 @@ class ModifyMilter(Milter.Base):
self.recipients.add("@".join(parse_addr(to)).lower()) self.recipients.add("@".join(parse_addr(to)).lower())
except Exception as e: except Exception as e:
self.logger.exception( self.logger.exception(
f"an exception occured in envrcpt function: {e}") f"an exception occured in envrcpt method: {e}")
return Milter.TEMPFAIL return Milter.TEMPFAIL
return Milter.CONTINUE return Milter.CONTINUE
@@ -231,64 +243,64 @@ class ModifyMilter(Milter.Base):
self.fields = None self.fields = None
self.fields_bytes = None self.fields_bytes = None
self.body_data = None self.body_data = None
needs = []
self._fp = BytesFeedParser(
_factory=MilterMessage, policy=default_policy)
self._keep_body = False
for rule in self.rules: for rule in self.rules:
needs += rule.needs() if rule.need_body():
self._keep_body = True
if "fields" in needs: break
self.fields = []
if "fields_bytes" in needs:
self.fields_bytes = []
if "body" in needs:
self.body_data = BytesIO()
except Exception as e: except Exception as e:
self.logger.exception( self.logger.exception(
f"an exception occured in data function: {e}") f"an exception occured in data method: {e}")
return Milter.TEMPFAIL return Milter.TEMPFAIL
return Milter.CONTINUE return Milter.CONTINUE
def header(self, name, value): def header(self, field, value):
try: try:
if self.fields_bytes is not None: # feed header line to BytesParser
self.fields_bytes.append( self._fp.feed(field + b": " + value + b"\r\n")
(name.encode("ascii", errors="surrogateescape"),
value.encode("ascii", errors="surrogateescape")))
if self.fields is not None: # remove surrogates from field and value
# remove surrogates from value field = remove_surrogates(field)
value = value.encode( value = remove_surrogates(value)
errors="surrogateescape").decode(errors="replace")
self.logger.debug(f"received header: {name}: {value}")
header = make_header(decode_header(value), errors="replace")
value = str(header).replace("\x00", "")
self.logger.debug(f"decoded header: {name}: {value}")
self.fields.append((name, value))
except Exception as e: except Exception as e:
self.logger.exception( self.logger.exception(
f"an exception occured in header function: {e}") f"an exception occured in header method: {e}")
return Milter.TEMPFAIL
return Milter.CONTINUE
def eoh(self):
try:
self._fp.feed(b"\r\n")
except Exception as e:
self.logger.exception(
f"an exception occured in eoh method: {e}")
return Milter.TEMPFAIL return Milter.TEMPFAIL
return Milter.CONTINUE return Milter.CONTINUE
def body(self, chunk): def body(self, chunk):
try: try:
if self.body_data is not None: if self._keep_body:
self.body_data.write(chunk) self._fp.feed(chunk)
except Exception as e: except Exception as e:
self.logger.exception( self.logger.exception(
f"an exception occured in body function: {e}") f"an exception occured in body method: {e}")
return Milter.TEMPFAIL return Milter.TEMPFAIL
return Milter.CONTINUE return Milter.CONTINUE
def eom(self): def eom(self):
try: try:
msg = self._fp.close()
for rule in self.rules: for rule in self.rules:
milter_action = rule.execute(self) milter_action = rule.execute(self, msg)
if milter_action is not None: if milter_action is not None:
if milter_action["action"] == "reject": if milter_action["action"] == "reject":
@@ -303,7 +315,7 @@ class ModifyMilter(Milter.Base):
except Exception as e: except Exception as e:
self.logger.exception( self.logger.exception(
f"an exception occured in eom function: {e}") f"an exception occured in eom method: {e}")
return Milter.TEMPFAIL return Milter.TEMPFAIL
return Milter.ACCEPT return Milter.ACCEPT

View File

@@ -21,10 +21,8 @@ from collections import defaultdict
from copy import copy from copy import copy
from datetime import datetime from datetime import datetime
from email.header import Header from email.header import Header
from email.parser import BytesFeedParser
from email.message import MIMEPart from email.message import MIMEPart
from email.policy import default as default_policy, SMTP from email.policy import SMTP
from shutil import copyfileobj
from pymodmilter import CustomLogger, Conditions from pymodmilter import CustomLogger, Conditions
@@ -37,7 +35,7 @@ def _replace_illegal_chars(string):
"\n", "") "\n", "")
def add_header(field, value, milter, idx=-1, pretend=False, def add_header(milter, msg, field, value, pretend=False,
logger=logging.getLogger(__name__)): logger=logging.getLogger(__name__)):
"""Add a mail header field.""" """Add a mail header field."""
header = f"{field}: {value}" header = f"{field}: {value}"
@@ -46,21 +44,18 @@ def add_header(field, value, milter, idx=-1, pretend=False,
else: else:
logger.info(f"add_header: {header[0:70]}") logger.info(f"add_header: {header[0:70]}")
if idx == -1: msg.add_header(field, value)
milter.fields.append((field, value))
else:
milter.fields.insert(idx, (field, value))
if pretend: if pretend:
return return
encoded_value = _replace_illegal_chars( encoded_value = _replace_illegal_chars(
Header(s=value).encode()) Header(s=value).encode())
milter.logger.debug(f"milter: addheader: {field}[{idx}]: {encoded_value}") milter.logger.debug(f"milter: addheader: {field}: {encoded_value}")
milter.addheader(field, encoded_value, idx) milter.addheader(field, encoded_value, -1)
def mod_header(field, value, milter, search=None, pretend=False, def mod_header(milter, msg, field, value, search=None, pretend=False,
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):
@@ -71,8 +66,9 @@ def mod_header(field, value, milter, search=None, pretend=False,
occ = defaultdict(int) occ = defaultdict(int)
for idx, (f, v) in enumerate(milter.fields): for i, (f, v) in enumerate(msg.items()):
occ[f] += 1 f_lower = f.lower()
occ[f_lower] += 1
if not field.match(f): if not field.match(f):
continue continue
@@ -93,12 +89,13 @@ def mod_header(field, value, milter, search=None, pretend=False,
header = f"{f}: {v}" header = f"{f}: {v}"
new_header = f"{f}: {new_v}" new_header = f"{f}: {new_v}"
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}")
else: else:
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.fields[idx] = (f, new_v) msg.replace_header(f, new_v, occ=occ[f_lower])
if pretend: if pretend:
continue continue
@@ -106,11 +103,11 @@ def mod_header(field, value, milter, search=None, pretend=False,
encoded_value = _replace_illegal_chars( encoded_value = _replace_illegal_chars(
Header(s=new_v).encode()) Header(s=new_v).encode())
milter.logger.debug( milter.logger.debug(
f"milter: chgheader: {f}[{occ[f]}]: {encoded_value}") f"milter: chgheader: {f}[{occ[f_lower]}]: {encoded_value}")
milter.chgheader(f, occ[f], encoded_value) milter.chgheader(f, occ[f_lower], encoded_value)
def del_header(field, milter, value=None, pretend=False, def del_header(milter, msg, 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):
@@ -119,14 +116,13 @@ def del_header(field, milter, value=None, pretend=False,
if isinstance(value, str): if isinstance(value, str):
value = re.compile(value, re.MULTILINE + re.DOTALL + re.IGNORECASE) value = re.compile(value, re.MULTILINE + re.DOTALL + re.IGNORECASE)
idx = -1
occ = defaultdict(int) occ = defaultdict(int)
# iterate a copy of milter.fields because elements may get removed # iterate a copy of milter.fields because elements may get removed
# during iteration # during iteration
for f, v in milter.fields.copy(): for f, v in msg.items():
idx += 1 f_lower = f.lower()
occ[f] += 1 occ[f_lower] += 1
if not field.match(f): if not field.match(f):
continue continue
@@ -140,16 +136,14 @@ def del_header(field, milter, value=None, pretend=False,
else: else:
logger.info(f"del_header: {header[0:70]}") logger.info(f"del_header: {header[0:70]}")
del milter.fields[idx] msg.remove_header(f, occ=occ[f_lower])
occ[f_lower] -= 1
if not pretend: if not pretend:
encoded_value = ""
milter.logger.debug( milter.logger.debug(
f"milter: chgheader: {f}[{occ[f]}]: {encoded_value}") f"milter: chgheader: {f}[{occ[f_lower]}]:")
milter.chgheader(f, occ[f], encoded_value) milter.chgheader(f, occ[f_lower], "")
idx -= 1
occ[f] -= 1
def _get_body_content(msg, body_type): def _get_body_content(msg, body_type):
@@ -273,31 +267,9 @@ def _inject_body(milter, msg):
return new_msg return new_msg
def add_disclaimer(text, html, action, policy, milter, pretend=False, def add_disclaimer(milter, msg, text, html, action, policy, pretend=False,
logger=logging.getLogger(__name__)): logger=logging.getLogger(__name__)):
"""Append or prepend a disclaimer to the mail body.""" """Append or prepend a disclaimer to the mail body."""
milter.body_data.seek(0)
fp = BytesFeedParser(policy=default_policy)
for field, value in milter.fields_bytes:
decoded_field = field.decode("ascii")
decoded_value = value.decode("ascii")
field_lower = decoded_field.lower()
if not field_lower.startswith("content-") and \
field_lower != "mime-version":
continue
logger.debug(
f"feed content header to message object: "
f"{decoded_field}: {decoded_value}")
fp.feed(field + b": " + value + b"\r\n")
fp.feed(b"\r\n")
logger.debug(f"feed body to message object")
fp.feed(milter.body_data.read())
logger.debug("parse message")
msg = fp.close()
update_headers = False update_headers = False
try: try:
@@ -367,48 +339,47 @@ def add_disclaimer(text, html, action, policy, milter, pretend=False,
"value": msg.get("Content-Transfer-Encoding"), "value": msg.get("Content-Transfer-Encoding"),
"modified": False}} "modified": False}}
for field, value in milter.fields.copy(): for field, value in msg.items():
field_lower = field.lower() field_lower = field.lower()
if field_lower in fields and fields[field_lower]["value"] is not None: if field_lower in fields and fields[field_lower]["value"] is not None:
mod_header(field=f"^{field}$", value=fields[field_lower]["value"], mod_header(milter, msg, field=f"^{field}$",
milter=milter, pretend=pretend, logger=logger) value=fields[field_lower]["value"],
pretend=pretend, logger=logger)
fields[field_lower]["modified"] = True fields[field_lower]["modified"] = True
elif field_lower.startswith("content-"): elif field_lower.startswith("content-"):
del_header(field=f"^{field}$", milter=milter, del_header(milter, msg, field=f"^{field}$",
pretend=pretend, logger=logger) pretend=pretend, logger=logger)
for field in fields.values(): for field in fields.values():
if not field["modified"] and field["value"] is not None: if not field["modified"] and field["value"] is not None:
add_header(field=field["field"], value=field["value"], add_header(milter, msg, field=field["field"], value=field["value"],
milter=milter, pretend=pretend, logger=logger) pretend=pretend, logger=logger)
def store(directory, milter, pretend=False, def store(milter, msg, directory, pretend=False,
logger=logging.getLogger(__name__)): logger=logging.getLogger(__name__)):
timestamp = datetime.now().strftime("%Y%m%d%H%M%S") timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
store_id = f"{timestamp}_{milter.qid}" store_id = f"{timestamp}_{milter.qid}"
datafile = os.path.join(directory, store_id) datafile = os.path.join(directory, store_id)
milter.body_data.seek(0) milter.body_data.seek(0)
logger.info(f"store message in file {datafile}") logger.info(f"store message in file {datafile}")
try: try:
with open(datafile, "wb") as fp: with open(datafile, "wb") as fp:
for field, value in milter.fields_bytes: fp.write(msg.as_bytes())
fp.write(field + b": " + value + b"\r\n")
copyfileobj(milter.body_data, fp)
except IOError as e: except IOError as e:
raise RuntimeError(f"unable to store message: {e}") raise RuntimeError(f"unable to store message: {e}")
class Action: class Action:
"""Action to implement a pre-configured action to perform on e-mails.""" """Action to implement a pre-configured action to perform on e-mails."""
_types = { _need_body_map = {
"add_header": ["fields"], "add_header": False,
"del_header": ["fields"], "del_header": False,
"mod_header": ["fields"], "mod_header": False,
"add_disclaimer": ["fields", "body"], "add_disclaimer": True,
"store": ["fields_bytes", "body"]} "store": True}
def __init__(self, name, local_addrs, conditions, action_type, args, def __init__(self, name, local_addrs, conditions, action_type, args,
loglevel=logging.INFO, pretend=False): loglevel=logging.INFO, pretend=False):
@@ -423,9 +394,9 @@ class Action:
self.pretend = pretend self.pretend = pretend
self._args = {} self._args = {}
if action_type not in self._types: if action_type not in self._need_body_map:
raise RuntimeError(f"invalid action type '{action_type}'") raise RuntimeError(f"invalid action type '{action_type}'")
self._needs = self._types[action_type] self._need_body = self._need_body_map[action_type]
try: try:
if action_type == "add_header": if action_type == "add_header":
@@ -498,16 +469,16 @@ class Action:
raise RuntimeError( raise RuntimeError(
f"mandatory argument not found: {e}") f"mandatory argument not found: {e}")
def needs(self): def need_body(self):
"""Return the needs of this action.""" """Return the needs of this action."""
return self._needs return self._need_body
def execute(self, milter, pretend=None): def execute(self, milter, msg, pretend=None):
"""Execute configured action.""" """Execute configured action."""
if pretend is None: if pretend is None:
pretend = self.pretend pretend = self.pretend
logger = CustomLogger(self.logger, {"qid": milter.qid}) logger = CustomLogger(self.logger, {"qid": milter.qid})
return self._func( return self._func(milter=milter, msg=msg, pretend=pretend,
milter=milter, pretend=pretend, logger=logger, **self._args) logger=logger, **self._args)