From d03d2b5f445657982509e4c97fa6bbd665665a73 Mon Sep 17 00:00:00 2001 From: Nathan Thorpe Date: Thu, 21 Apr 2022 17:03:54 -0700 Subject: [PATCH] pep8 and tests fix --- parsedmarc/__init__.py | 17 +++++++++++------ parsedmarc/cli.py | 33 +++++++++++++++++---------------- parsedmarc/mail/__init__.py | 5 +++++ parsedmarc/mail/gmail.py | 2 +- 4 files changed, 34 insertions(+), 23 deletions(-) diff --git a/parsedmarc/__init__.py b/parsedmarc/__init__.py index 9245e41..c2ab46a 100644 --- a/parsedmarc/__init__.py +++ b/parsedmarc/__init__.py @@ -1034,10 +1034,11 @@ def get_dmarc_reports_from_mailbox(connection: MailboxConnection, nameservers (list): A list of DNS nameservers to query dns_timeout (float): Set the DNS query timeout strip_attachment_payloads (bool): Remove attachment payloads from - forensic report results + forensic report results results (dict): Results from the previous run batch_size (int): Number of messages to read and process before saving - create_folders (bool): Whether to create the destination folders (not used in watch) + create_folders (bool): Whether to create the destination folders + (not used in watch) Returns: OrderedDict: Lists of ``aggregate_reports`` and ``forensic_reports`` @@ -1085,12 +1086,13 @@ def get_dmarc_reports_from_mailbox(connection: MailboxConnection, )) msg_content = connection.fetch_message(msg_uid) try: + sa = strip_attachment_payloads parsed_email = parse_report_email(msg_content, nameservers=nameservers, dns_timeout=dns_timeout, ip_db_path=ip_db_path, offline=offline, - strip_attachment_payloads=strip_attachment_payloads, + strip_attachment_payloads=sa, keep_alive=connection.keepalive) if parsed_email["report_type"] == "aggregate": aggregate_reports.append(parsed_email["report"]) @@ -1202,7 +1204,8 @@ def watch_inbox(mailbox_connection: MailboxConnection, dns_timeout=6.0, strip_attachment_payloads=False, batch_size=None): """ - Watches the mailbox for new messages and sends the results to a callback function + Watches the mailbox for new messages and + sends the results to a callback function Args: mailbox_connection: The mailbox connection object callback: The callback function to receive the parsing results @@ -1223,6 +1226,7 @@ def watch_inbox(mailbox_connection: MailboxConnection, """ def check_callback(connection): + sa = strip_attachment_payloads res = get_dmarc_reports_from_mailbox(connection=connection, reports_folder=reports_folder, archive_folder=archive_folder, @@ -1232,12 +1236,13 @@ def watch_inbox(mailbox_connection: MailboxConnection, offline=offline, nameservers=nameservers, dns_timeout=dns_timeout, - strip_attachment_payloads=strip_attachment_payloads, + strip_attachment_payloads=sa, batch_size=batch_size, create_folders=False) callback(res) - mailbox_connection.watch(check_callback=check_callback, check_timeout=check_timeout) + mailbox_connection.watch(check_callback=check_callback, + check_timeout=check_timeout) def save_output(results, output_directory="output", diff --git a/parsedmarc/cli.py b/parsedmarc/cli.py index 7ba3a37..99391ae 100644 --- a/parsedmarc/cli.py +++ b/parsedmarc/cli.py @@ -240,7 +240,7 @@ def _main(): args = arg_parser.parse_args() - gmail_api_scopes = ['https://www.googleapis.com/auth/gmail.modify'] + default_gmail_api_scope = 'https://www.googleapis.com/auth/gmail.modify' opts = Namespace(file_path=args.file_path, config_file=args.config_file, @@ -315,11 +315,11 @@ def _main(): gmail_api_credentials_file=None, gmail_api_token_file=None, gmail_api_include_spam_trash=False, - gmail_api_scopes=gmail_api_scopes, + gmail_api_scopes=[], log_file=args.log_file, n_procs=1, chunk_size=1, - ip_db_path = None + ip_db_path=None ) args = arg_parser.parse_args() @@ -628,19 +628,18 @@ def _main(): opts.syslog_port = syslog_config["port"] else: opts.syslog_port = 514 - + if "gmail_api" in config.sections(): gmail_api_config = config["gmail_api"] - opts.gmail_api_credentials_file = gmail_api_config.get("credentials_file", None) - opts.gmail_api_token_file = gmail_api_config.get("token_file", ".token") - opts.gmail_api_reports_label = gmail_api_config.get("reports_label", "INBOX") - opts.gmail_api_archive_label = gmail_api_config.get("archive_label", "DMARC Archive") - opts.gmail_api_include_spam_trash = gmail_api_config.getboolean("include_spam_trash", False) - opts.gmail_api_scopes = str.split(gmail_api_config.get("scopes", - "https://www.googleapis.com/auth/gmail.modify"), - ",") - opts.gmail_api_delete = gmail_api_config.getboolean("delete", None) - opts.gmail_api_test = gmail_api_config.getboolean("test", False) + opts.gmail_api_credentials_file = \ + gmail_api_config.get("credentials_file") + opts.gmail_api_token_file = \ + gmail_api_config.get("token_file", ".token") + opts.gmail_api_include_spam_trash = \ + gmail_api_config.getboolean("include_spam_trash", False) + opts.gmail_api_scopes = \ + gmail_api_config.get("scopes", + default_gmail_api_scope).split(',') logger.setLevel(logging.WARNING) @@ -813,8 +812,10 @@ def _main(): if opts.gmail_api_credentials_file: if opts.gmail_api_delete: if 'https://mail.google.com/' not in opts.gmail_api_scopes: - logger.error("Message deletion requires scope 'https://mail.google.com/'. " - "Add the scope and remove token file to acquire proper access.") + logger.error("Message deletion requires scope" + " 'https://mail.google.com/'. " + "Add the scope and remove token file " + "to acquire proper access.") opts.gmail_api_delete = False try: diff --git a/parsedmarc/mail/__init__.py b/parsedmarc/mail/__init__.py index c376024..df3c4f2 100644 --- a/parsedmarc/mail/__init__.py +++ b/parsedmarc/mail/__init__.py @@ -2,3 +2,8 @@ from parsedmarc.mail.mailbox_connection import MailboxConnection from parsedmarc.mail.graph import MSGraphConnection from parsedmarc.mail.gmail import GmailConnection from parsedmarc.mail.imap import IMAPConnection + +__all__ = ["MailboxConnection", + "MSGraphConnection", + "GmailConnection", + "IMAPConnection"] diff --git a/parsedmarc/mail/gmail.py b/parsedmarc/mail/gmail.py index 51491d5..eaa40ef 100644 --- a/parsedmarc/mail/gmail.py +++ b/parsedmarc/mail/gmail.py @@ -9,7 +9,7 @@ from google.oauth2.credentials import Credentials from google_auth_oauthlib.flow import InstalledAppFlow from googleapiclient.discovery import build -from parsedmarc import MailboxConnection +from parsedmarc.mail.mailbox_connection import MailboxConnection logger = logging.getLogger("parsedmarc")