pep8 and tests fix

This commit is contained in:
Nathan Thorpe
2022-04-21 17:03:54 -07:00
parent cf682337e9
commit d03d2b5f44
4 changed files with 34 additions and 23 deletions

View File

@@ -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",

View File

@@ -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:

View File

@@ -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"]

View File

@@ -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")