From bcf0acef349fd269da266c4e68cc68e2c2ee9a88 Mon Sep 17 00:00:00 2001
From: Sean Whalen <whalenster@gmail.com>
Date: Sun, 25 Mar 2018 23:06:34 -0400
Subject: [PATCH] Fix exception handling

---
 parsedmarc/cli.py     | 34 +++++++++++++++++++---------------
 parsedmarc/elastic.py | 17 +++++++++--------
 requirements.txt      |  1 +
 setup.py              |  3 +--
 4 files changed, 30 insertions(+), 25 deletions(-)

diff --git a/parsedmarc/cli.py b/parsedmarc/cli.py
index e26789f..7f5d7a5 100644
--- a/parsedmarc/cli.py
+++ b/parsedmarc/cli.py
@@ -18,23 +18,28 @@ from parsedmarc import logger, IMAPError, get_dmarc_reports_from_inbox, \
 
 
 def _main():
-    """Called when the module in executed"""
+    """Called when the module is executed"""
     def process_reports(reports_):
-        try:
-            print(json.dumps(reports_, ensure_ascii=False, indent=2), "\n")
-            if args.save_aggregate:
-                for report in reports_["aggregate_reports"]:
+        print(json.dumps(reports_, ensure_ascii=False, indent=2), "\n")
+        if args.save_aggregate:
+            for report in reports_["aggregate_reports"]:
+                try:
                     elastic.save_aggregate_report_to_elasticsearch(report)
-                    sleep(1)
-            if args.save_forensic:
-                for report in reports_["forensic_reports"]:
+                except elastic.AlreadySaved as warning:
+                    logger.warning(warning.__str__())
+                except ElasticsearchException as error_:
+                    logger.error("Elasticsearch Error: {0}".format(
+                        error_.__str__()))
+                    exit(1)
+        if args.save_forensic:
+            for report in reports_["forensic_reports"]:
+                try:
                     elastic.save_forensic_report_to_elasticsearch(report)
-                    sleep(1)
-        except elastic.AlreadySaved as error_:
-            logger.warning(error_.__str__())
-        except ElasticsearchException as error_:
-            logger.error("Elasticsearch Error: {0}".format(error_.__str__()))
-            exit(1)
+                except elastic.AlreadySaved as warning:
+                    logger.warning(warning.__str__())
+                except ElasticsearchException as error_:
+                    logger.error("Elasticsearch Error: {0}".format(
+                        error_.__str__()))
 
     arg_parser = ArgumentParser(description="Parses DMARC reports")
     arg_parser.add_argument("file_path", nargs="*",
@@ -92,7 +97,6 @@ def _main():
                             help="Email the results using this filename")
     arg_parser.add_argument("-M", "--outgoing-message",
                             help="Email the results using this message")
-
     arg_parser.add_argument("-i", "--idle", action="store_true",
                             help="Use an IMAP IDLE connection to process "
                                  "reports as they arrive in the inbox")
diff --git a/parsedmarc/elastic.py b/parsedmarc/elastic.py
index e3b500f..a9f6355 100644
--- a/parsedmarc/elastic.py
+++ b/parsedmarc/elastic.py
@@ -83,6 +83,7 @@ class AggregateReportDoc(DocType):
     def save(self, ** kwargs):
         self.passed_dmarc = False
         self.passed_dmarc = self.spf_aligned or self.dkim_aligned
+
         return super().save(** kwargs)
 
 
@@ -155,7 +156,7 @@ class ForensicReportDoc(DocType):
     sample = Object(ForensicSampleDoc)
 
 
-class AlreadySaved(RuntimeError):
+class AlreadySaved(ValueError):
     """Raised when a report to be saved matches an existing report"""
 
 
@@ -187,8 +188,7 @@ def save_aggregate_report_to_elasticsearch(aggregate_report):
         aggregate_report (OrderedDict): A parsed forensic report
 
     Raises:
-        AlreadySaved
-
+            AlreadySaved
     """
     aggregate_report = aggregate_report.copy()
     metadata = aggregate_report["report_metadata"]
@@ -209,10 +209,10 @@ def save_aggregate_report_to_elasticsearch(aggregate_report):
                                                                     org_name,
                                                                     domain))
 
-    aggregate_report["begin_date"] = parsedmarc.human_timestamp_to_datetime(
-        metadata["begin_date"])
-    aggregate_report["end_date"] = parsedmarc.human_timestamp_to_datetime(
-        metadata["end_date"])
+    begin_date = parsedmarc.human_timestamp_to_datetime(metadata["begin_date"])
+    end_date = parsedmarc.human_timestamp_to_datetime(metadata["end_date"])
+    aggregate_report["begin_date"] = begin_date
+    aggregate_report["end_date"] = end_date
     date_range = (aggregate_report["begin_date"],
                   aggregate_report["end_date"])
     published_policy = PublishedPolicy(
@@ -290,7 +290,8 @@ def save_forensic_report_to_elasticsearch(forensic_report):
     to_query = {"match": {"sample.headers.to": headers["to"]}}
     from_query = {"match": {"sample.headers.from": headers["from"]}}
     subject_query = {"match": {"sample.headers.subject": headers["subject"]}}
-    search.query = Q(to_query) & Q(from_query) & Q(subject_query)
+    arrival_date_query = {"match": {"sample.headers.arrival_date": forensic_report["arrival_date_utc"]}}
+    search.query = Q(to_query) & Q(from_query) & Q(subject_query) & Q(arrival_date_query)
     existing = search.execute()
 
     if len(existing) > 0:
diff --git a/requirements.txt b/requirements.txt
index 2fe92e3..ee86a58 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,6 +7,7 @@ dnspython
 imapclient
 mail-parser
 dateparser
+elasticsearch
 elasticsearch-dsl
 flake8
 sphinx
diff --git a/setup.py b/setup.py
index 772afc4..74ac4e8 100644
--- a/setup.py
+++ b/setup.py
@@ -7,7 +7,6 @@ https://packaging.python.org/en/latest/distributing.html
 https://github.com/pypa/sampleproject
 """
 
-from __future__ import absolute_import
 
 # Always prefer setuptools over distutils
 from setuptools import setup
@@ -93,7 +92,7 @@ setup(
     # https://packaging.python.org/en/latest/requirements.html
     install_requires=['dnspython', 'publicsuffix', 'xmltodict', 'geoip2',
                       'dnspython', 'imapclient', 'mail-parser', 'dateparser',
-                      'elasticsearch-dsl'
+                      'elasticsearch', 'elasticsearch-dsl'
                       ],
 
     entry_points={