From 4a7df9804b13091e1da58eadc3b7f84ec4a89606 Mon Sep 17 00:00:00 2001
From: Sean Whalen <whalenster@gmail.com>
Date: Mon, 19 Mar 2018 12:44:31 -0400
Subject: [PATCH] Add forgotten code

---
 parsedmarc/cli.py     | 12 ++++++++++++
 parsedmarc/elastic.py | 10 +++++-----
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/parsedmarc/cli.py b/parsedmarc/cli.py
index 0232082..d0ba4c8 100644
--- a/parsedmarc/cli.py
+++ b/parsedmarc/cli.py
@@ -21,9 +21,21 @@ def _main():
 
         Args:
             results_ (OrderedDict): Parsing results
+            save_aggregate (bool): Save Aggregate reports to Elasticsearch
+            save_forensic (bool): Save forensic reports to Elasticsearch
         """
         print(json.dumps(results_, ensure_ascii=False, indent=2), "\n")
 
+        try:
+            if save_aggregate:
+                for report in results["aggregate_reports"]:
+                    elastic.save_aggregate_report_to_elasticsearch(report)
+            if save_forensic:
+                for report in results["forensic_reports"]:
+                    elastic.save_forensic_report_to_elasticsearch(report)
+        except elastic.AlreadySaved as exception:
+            logger.warning(exception.__str__())
+
     arg_parser = ArgumentParser(description="Parses DMARC reports")
     arg_parser.add_argument("file_path", nargs="*",
                             help="one or more paths to aggregate or forensic "
diff --git a/parsedmarc/elastic.py b/parsedmarc/elastic.py
index 2ec913d..4b00a10 100644
--- a/parsedmarc/elastic.py
+++ b/parsedmarc/elastic.py
@@ -154,7 +154,7 @@ class ForensicReportDoc(DocType):
     sample = Object(ForensicSampleDoc)
 
 
-class ExistingReport(RuntimeError):
+class AlreadySaved(RuntimeError):
     """Raised when a report to be saved matches an existing report"""
 
 
@@ -186,7 +186,7 @@ def save_aggregate_report_to_elasticsearch(aggregate_report):
         aggregate_report (OrderedDict): A parsed forensic report
 
     Raises:
-        ExistingReport
+        AlreadySaved
 
     """
     metadata = aggregate_report["report_metadata"]
@@ -202,7 +202,7 @@ def save_aggregate_report_to_elasticsearch(aggregate_report):
     search.query = org_name_query & report_id_query & domain_query
     existing = search.execute()
     if len(existing) > 0:
-        raise ExistingReport("A matching aggregate report already exists")
+        raise AlreadySaved("A matching aggregate report already exists")
 
     aggregate_report["begin_date"] = parsedmarc.human_timestamp_to_datetime(
         metadata["begin_date"])
@@ -267,7 +267,7 @@ def save_forensic_report_to_elasticsearch(forensic_report):
             forensic_report (OrderedDict): A parsed forensic report
 
         Raises:
-            ExistingReport
+            AlreadySaved
 
         """
     sample_date = forensic_report["parsed_sample"]["date"]
@@ -288,7 +288,7 @@ def save_forensic_report_to_elasticsearch(forensic_report):
     existing = search.execute()
 
     if len(existing) > 0:
-        raise ExistingReport(" A matching forensic report already exists")
+        raise AlreadySaved(" A matching forensic report already exists")
 
     parsed_sample = forensic_report["parsed_sample"]
     sample = ForensicSampleDoc(