From 3615ad3799f058e8012b39f708aa500f51aa2cf4 Mon Sep 17 00:00:00 2001 From: Silvian I Date: Mon, 31 May 2021 15:40:57 +0200 Subject: [PATCH] Fix server connection timeout while processiong large dmarc files --- parsedmarc/__init__.py | 37 ++++++++++++++++++++++++++----------- 1 file changed, 26 insertions(+), 11 deletions(-) diff --git a/parsedmarc/__init__.py b/parsedmarc/__init__.py index 562a0a8..9de72c1 100644 --- a/parsedmarc/__init__.py +++ b/parsedmarc/__init__.py @@ -203,7 +203,7 @@ def _parse_report_record(record, offline=False, nameservers=None, def parse_aggregate_report_xml(xml, offline=False, nameservers=None, - timeout=2.0, parallel=False): + timeout=2.0, parallel=False, server=None): """Parses a DMARC XML report string and returns a consistent OrderedDict Args: @@ -213,6 +213,7 @@ def parse_aggregate_report_xml(xml, offline=False, nameservers=None, (Cloudflare's public DNS resolvers by default) timeout (float): Sets the DNS timeout in seconds parallel (bool): Parallel processing + server (IMAPClient): Connection object Returns: OrderedDict: The parsed aggregate DMARC report @@ -304,8 +305,13 @@ def parse_aggregate_report_xml(xml, offline=False, nameservers=None, new_report["policy_published"] = new_policy_published if type(report["record"]) == list: - for record in report["record"]: - report_record = _parse_report_record(record, + for i in range(len(report["record"])): + if server is not None and i > 0 and i % 20 == 0: + logger.debug("Sending noop cmd") + server.noop() + logger.debug("Processed {0}/{1}".format( + i, len(report["record"]))) + report_record = _parse_report_record(report["record"][i], offline=offline, nameservers=nameservers, dns_timeout=timeout, @@ -385,7 +391,8 @@ def extract_xml(input_): def parse_aggregate_report_file(_input, offline=False, nameservers=None, dns_timeout=2.0, - parallel=False): + parallel=False, + server=None): """Parses a file at the given path, a file-like object. or bytes as a aggregate DMARC report @@ -396,6 +403,7 @@ def parse_aggregate_report_file(_input, offline=False, nameservers=None, (Cloudflare's public DNS resolvers by default) dns_timeout (float): Sets the DNS timeout in seconds parallel (bool): Parallel processing + server (IMAPClient): Connection object Returns: OrderedDict: The parsed DMARC aggregate report @@ -406,7 +414,8 @@ def parse_aggregate_report_file(_input, offline=False, nameservers=None, offline=offline, nameservers=nameservers, timeout=dns_timeout, - parallel=parallel) + parallel=parallel, + server=server) def parsed_aggregate_reports_to_csv_rows(reports): @@ -738,7 +747,7 @@ def parsed_forensic_reports_to_csv(reports): def parse_report_email(input_, offline=False, nameservers=None, dns_timeout=2.0, strip_attachment_payloads=False, - parallel=False): + parallel=False, server=None): """ Parses a DMARC report from an email @@ -750,6 +759,7 @@ def parse_report_email(input_, offline=False, nameservers=None, strip_attachment_payloads (bool): Remove attachment payloads from forensic report results parallel (bool): Parallel processing + server (IMAPClient): Connection object Returns: OrderedDict: @@ -813,7 +823,8 @@ def parse_report_email(input_, offline=False, nameservers=None, offline=offline, nameservers=ns, dns_timeout=dns_timeout, - parallel=parallel) + parallel=parallel, + server=server) result = OrderedDict([("report_type", "aggregate"), ("report", aggregate_report)]) return result @@ -863,7 +874,7 @@ def parse_report_email(input_, offline=False, nameservers=None, def parse_report_file(input_, nameservers=None, dns_timeout=2.0, strip_attachment_payloads=False, - offline=False, parallel=False): + offline=False, parallel=False, server=None): """Parses a DMARC aggregate or forensic file at the given path, a file-like object. or bytes @@ -876,6 +887,7 @@ def parse_report_file(input_, nameservers=None, dns_timeout=2.0, forensic report results offline (bool): Do not make online queries for geolocation or DNS parallel (bool): Parallel processing + server (IMAPClient): Connection object Returns: OrderedDict: The parsed DMARC report @@ -895,7 +907,8 @@ def parse_report_file(input_, nameservers=None, dns_timeout=2.0, offline=offline, nameservers=nameservers, dns_timeout=dns_timeout, - parallel=parallel) + parallel=parallel, + server=server) results = OrderedDict([("report_type", "aggregate"), ("report", report)]) except InvalidAggregateReport: @@ -906,7 +919,8 @@ def parse_report_file(input_, nameservers=None, dns_timeout=2.0, nameservers=nameservers, dns_timeout=dns_timeout, strip_attachment_payloads=sa, - parallel=parallel) + parallel=parallel, + server=server) except InvalidDMARCReport: raise InvalidDMARCReport("Not a valid aggregate or forensic " "report") @@ -1083,7 +1097,8 @@ def get_dmarc_reports_from_inbox(connection=None, nameservers=nameservers, dns_timeout=dns_timeout, offline=offline, - strip_attachment_payloads=sa) + strip_attachment_payloads=sa, + server=server) if parsed_email["report_type"] == "aggregate": aggregate_reports.append(parsed_email["report"]) aggregate_report_msg_uids.append(msg_uid)