From af8eb8c141a12fc1c2e93c3dad62bb8aee733e9c Mon Sep 17 00:00:00 2001 From: Thomas Oettli Date: Wed, 6 Oct 2021 20:58:04 +0200 Subject: [PATCH] improve message serialization and make cli.py PEP8 comform --- pyquarantine/__init__.py | 20 ++++++++- pyquarantine/cli.py | 94 ++++++++++++++++++++-------------------- pyquarantine/modify.py | 2 +- pyquarantine/storage.py | 2 +- 4 files changed, 68 insertions(+), 50 deletions(-) diff --git a/pyquarantine/__init__.py b/pyquarantine/__init__.py index 484855f..bfb7886 100644 --- a/pyquarantine/__init__.py +++ b/pyquarantine/__init__.py @@ -93,12 +93,28 @@ class QuarantineMilter(Milter.Base): self.logger.debug(f"delheader: {field}[{idx}]") super().chgheader(field, idx, value) + def msg_as_bytes(self): + try: + data = self.msg.as_bytes() + except Exception: + self.logger.warning( + "unable to serialize message as bytes, " + "try to serialize as str and encode") + try: + data = self.msg.as_string().encode("ascii", errors="replace") + except Exception as e: + self.logger.exception( + "unable to serialize message, giving up") + raise e + + return data + def update_headers(self, old_headers): if self.msg.is_multipart() and not self.msg["MIME-Version"]: self.msg.add_header("MIME-Version", "1.0") # serialize the message object so it updates its internal strucure - self.msg.as_bytes() + self.msg_as_bytes() old_headers = [(f, f.lower(), v) for f, v in old_headers] headers = [(f, f.lower(), v) for f, v in self.msg.items()] @@ -121,7 +137,7 @@ class QuarantineMilter(Milter.Base): def _replacebody(self): if not self._body_changed: return - data = self.msg.as_bytes() + 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:]) diff --git a/pyquarantine/cli.py b/pyquarantine/cli.py index 461d5f5..a96a987 100644 --- a/pyquarantine/cli.py +++ b/pyquarantine/cli.py @@ -194,11 +194,11 @@ def list_whitelist(quarantines, args): return # transform some values to strings - for entry_id, entry in entries.items(): - entries[entry_id]["permanent_str"] = str(entry["permanent"]) - entries[entry_id]["created_str"] = entry["created"].strftime( + for eid, entry in entries.items(): + entries[eid]["permanent_str"] = str(entry["permanent"]) + entries[eid]["created_str"] = entry["created"].strftime( '%Y-%m-%d %H:%M:%S') - entries[entry_id]["last_used_str"] = entry["last_used"].strftime( + entries[eid]["last_used_str"] = entry["last_used"].strftime( '%Y-%m-%d %H:%M:%S') print_table( @@ -227,11 +227,11 @@ def add_whitelist_entry(quarantines, args): if not args.force: # the entry is already covered by others - for entry_id, entry in entries.items(): - entries[entry_id]["permanent_str"] = str(entry["permanent"]) - entries[entry_id]["created_str"] = entry["created"].strftime( + for eid, entry in entries.items(): + entries[eid]["permanent_str"] = str(entry["permanent"]) + entries[eid]["created_str"] = entry["created"].strftime( '%Y-%m-%d %H:%M:%S') - entries[entry_id]["last_used_str"] = entry["last_used"].strftime( + entries[eid]["last_used_str"] = entry["last_used"].strftime( '%Y-%m-%d %H:%M:%S') print_table( [ @@ -342,144 +342,145 @@ def main(): list_parser.set_defaults(func=list_quarantines) # quarantine command group - quarantine_parser = subparsers.add_parser( + quar_parser = subparsers.add_parser( "quarantine", description="Manage quarantines.", help="Manage quarantines.", formatter_class=formatter_class) - quarantine_parser.add_argument( + quar_parser.add_argument( "quarantine", metavar="QUARANTINE", help="Quarantine name.") - quarantine_subparsers = quarantine_parser.add_subparsers( + quar_subparsers = quar_parser.add_subparsers( dest="command", title="Quarantine commands") - quarantine_subparsers.required = True + quar_subparsers.required = True # quarantine list command - quarantine_list_parser = quarantine_subparsers.add_parser( + quar_list_parser = quar_subparsers.add_parser( "list", description="List emails in quarantines.", help="List emails in quarantine.", formatter_class=formatter_class) - quarantine_list_parser.add_argument( + quar_list_parser.add_argument( "-f", "--from", dest="mailfrom", help="Filter emails by from address.", default=None, nargs="+") - quarantine_list_parser.add_argument( + quar_list_parser.add_argument( "-t", "--to", dest="recipients", help="Filter emails by recipient address.", default=None, nargs="+") - quarantine_list_parser.add_argument( + quar_list_parser.add_argument( "-o", "--older-than", dest="older_than", help="Filter emails by age (days).", default=None, type=float) - quarantine_list_parser.add_argument( + quar_list_parser.add_argument( "-b", "--batch", - help="Print results using only email quarantine IDs, each on a new line.", + help="Print results using only email quarantine IDs, " + "each on a new line.", action="store_true") - quarantine_list_parser.set_defaults(func=list_quarantine_emails) + quar_list_parser.set_defaults(func=list_quarantine_emails) # quarantine notify command - quarantine_notify_parser = quarantine_subparsers.add_parser( + quar_notify_parser = quar_subparsers.add_parser( "notify", description="Notify recipient about email in quarantine.", help="Notify recipient about email in quarantine.", formatter_class=formatter_class) - quarantine_notify_parser.add_argument( + quar_notify_parser.add_argument( "quarantine_id", metavar="ID", help="Quarantine ID.") - quarantine_notify_parser_group = quarantine_notify_parser.add_mutually_exclusive_group( + quar_notify_parser_grp = quar_notify_parser.add_mutually_exclusive_group( required=True) - quarantine_notify_parser_group.add_argument( + quar_notify_parser_grp.add_argument( "-t", "--to", dest="recipient", help="Release email for one recipient address.") - quarantine_notify_parser_group.add_argument( + quar_notify_parser_grp.add_argument( "-a", "--all", help="Release email for all recipients.", action="store_true") - quarantine_notify_parser.set_defaults(func=notify) + quar_notify_parser.set_defaults(func=notify) # quarantine release command - quarantine_release_parser = quarantine_subparsers.add_parser( + quar_release_parser = quar_subparsers.add_parser( "release", description="Release email from quarantine.", help="Release email from quarantine.", formatter_class=formatter_class) - quarantine_release_parser.add_argument( + quar_release_parser.add_argument( "quarantine_id", metavar="ID", help="Quarantine ID.") - quarantine_release_parser.add_argument( + quar_release_parser.add_argument( "-n", "--disable-syslog", dest="syslog", help="Disable syslog messages.", action="store_false") - quarantine_release_parser_group = quarantine_release_parser.add_mutually_exclusive_group( + quar_release_parser_grp = quar_release_parser.add_mutually_exclusive_group( required=True) - quarantine_release_parser_group.add_argument( + quar_release_parser_grp.add_argument( "-t", "--to", dest="recipient", help="Release email for one recipient address.") - quarantine_release_parser_group.add_argument( + quar_release_parser_grp.add_argument( "-a", "--all", help="Release email for all recipients.", action="store_true") - quarantine_release_parser.set_defaults(func=release) + quar_release_parser.set_defaults(func=release) # quarantine delete command - quarantine_delete_parser = quarantine_subparsers.add_parser( + quar_delete_parser = quar_subparsers.add_parser( "delete", description="Delete email from quarantine.", help="Delete email from quarantine.", formatter_class=formatter_class) - quarantine_delete_parser.add_argument( + quar_delete_parser.add_argument( "quarantine_id", metavar="ID", help="Quarantine ID.") - quarantine_delete_parser.add_argument( + quar_delete_parser.add_argument( "-n", "--disable-syslog", dest="syslog", help="Disable syslog messages.", action="store_false") - quarantine_delete_parser_group = quarantine_delete_parser.add_mutually_exclusive_group( + quar_delete_parser_grp = quar_delete_parser.add_mutually_exclusive_group( required=True) - quarantine_delete_parser_group.add_argument( + quar_delete_parser_grp.add_argument( "-t", "--to", dest="recipient", help="Delete email for one recipient address.") - quarantine_delete_parser_group.add_argument( + quar_delete_parser_grp.add_argument( "-a", "--all", help="Delete email for all recipients.", action="store_true") - quarantine_delete_parser.set_defaults(func=delete) + quar_delete_parser.set_defaults(func=delete) # quarantine get command - quarantine_get_parser = quarantine_subparsers.add_parser( + quar_get_parser = quar_subparsers.add_parser( "get", description="Get email from quarantine.", help="Get email from quarantine", formatter_class=formatter_class) - quarantine_get_parser.add_argument( + quar_get_parser.add_argument( "quarantine_id", metavar="ID", help="Quarantine ID.") - quarantine_get_parser.set_defaults(func=get) + quar_get_parser.set_defaults(func=get) # quarantine metadata command - quarantine_metadata_parser = quarantine_subparsers.add_parser( + quar_metadata_parser = quar_subparsers.add_parser( "metadata", description="Get metadata of email from quarantine.", help="Get metadata of email from quarantine", formatter_class=formatter_class) - quarantine_metadata_parser.add_argument( + quar_metadata_parser.add_argument( "quarantine_id", metavar="ID", help="Quarantine ID.") - quarantine_metadata_parser.set_defaults(func=metadata) + quar_metadata_parser.set_defaults(func=metadata) # whitelist command group whitelist_parser = subparsers.add_parser( @@ -546,7 +547,8 @@ def main(): action="store_true") whitelist_add_parser.add_argument( "--force", - help="Force adding an entry, even if already covered by another entry.", + help="Force adding an entry, " + "even if already covered by another entry.", action="store_true") whitelist_add_parser.set_defaults(func=add_whitelist_entry) # whitelist delete command diff --git a/pyquarantine/modify.py b/pyquarantine/modify.py index 070a7dc..1be04e8 100644 --- a/pyquarantine/modify.py +++ b/pyquarantine/modify.py @@ -204,7 +204,7 @@ def _inject_body(milter): def _wrap_message(milter): attachment = MIMEPart(policy=SMTPUTF8) - attachment.set_content(milter.msg.as_bytes(), + attachment.set_content(milter.msg_as_bytes(), maintype="plain", subtype="text", disposition="attachment", filename=f"{milter.qid}.eml", diff --git a/pyquarantine/storage.py b/pyquarantine/storage.py index c696513..727a92b 100644 --- a/pyquarantine/storage.py +++ b/pyquarantine/storage.py @@ -212,7 +212,7 @@ class FileMailStorage(BaseMailStorage): except Exception: subject = "" else: - data = milter.msg.as_bytes + data = milter.msg_as_bytes mailfrom = milter.msginfo["mailfrom"] recipients = list(milter.msginfo["rcpts"]) subject = milter.msg["subject"] or ""