From 419e8a68b20167a4fe42701c95c16c3a77c7ae76 Mon Sep 17 00:00:00 2001
From: Sean Whalen <whalenster@gmail.com>
Date: Wed, 9 Jan 2019 14:35:12 -0500
Subject: [PATCH] Add options for monthly Elasticsearch indexes (issue #41)

---
 parsedmarc/cli.py     | 13 ++++++++++---
 parsedmarc/elastic.py | 19 +++++++++++++++----
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/parsedmarc/cli.py b/parsedmarc/cli.py
index bd6e496..430bf9f 100644
--- a/parsedmarc/cli.py
+++ b/parsedmarc/cli.py
@@ -42,7 +42,8 @@ def _main():
                     if args.elasticsearch_host:
                         elastic.save_aggregate_report_to_elasticsearch(
                             report,
-                            index_suffix=args.elasticsearch_index_suffix)
+                            index_suffix=args.elasticsearch_index_suffix,
+                            monthly_indexes=args.elasticsearch_monthly_indexes)
                 except elastic.AlreadySaved as warning:
                     logger.warning(warning.__str__())
                 except elastic.ElasticsearchError as error_:
@@ -69,7 +70,8 @@ def _main():
                     if args.elasticsearch_host:
                         elastic.save_forensic_report_to_elasticsearch(
                             report,
-                            index_suffix=args.elasticsearch_index_suffix)
+                            index_suffix=args.elasticsearch_index_suffix,
+                            monthly_indexes=args.elasticsearch_monthly_indexes)
                 except elastic.AlreadySaved as warning:
                     logger.warning(warning.__str__())
                 except elastic.ElasticsearchError as error_:
@@ -146,6 +148,10 @@ def _main():
                             help="Use SSL when connecting to Elasticsearch")
     arg_parser.add_argument("--elasticsearch-ssl-cert-path", default=None,
                             help="Path to the Elasticsearch SSL certificate")
+    arg_parser.add_argument("--elasticsearch-monthly-indexes",
+                            action="store_true", default=False,
+                            help="Use monthly Elasticsearch indexes instead "
+                                 "of daily indexes")
     arg_parser.add_argument("--hec", help="the URL to a Splunk HTTP Event "
                                           "Collector (HEC)")
     arg_parser.add_argument("--hec-token", help="the authorization token for "
@@ -249,7 +255,8 @@ def _main():
                         es_aggregate_index, suffix)
                     es_forensic_index = "{0}_{1}".format(
                         es_forensic_index, suffix)
-                elastic.set_hosts(args.elasticsearch_host, args.elasticsearch_use_ssl,
+                elastic.set_hosts(args.elasticsearch_host,
+                                  args.elasticsearch_use_ssl,
                                   args.elasticsearch_ssl_cert_path)
                 elastic.migrate_indexes(aggregate_indexes=[es_aggregate_index],
                                         forensic_indexes=[es_forensic_index])
diff --git a/parsedmarc/elastic.py b/parsedmarc/elastic.py
index 2199005..72db2ec 100644
--- a/parsedmarc/elastic.py
+++ b/parsedmarc/elastic.py
@@ -265,13 +265,15 @@ def migrate_indexes(aggregate_indexes=None, forensic_indexes=None):
 
 
 def save_aggregate_report_to_elasticsearch(aggregate_report,
-                                           index_suffix=None):
+                                           index_suffix=None,
+                                           monthly_indexes=False):
     """
     Saves a parsed DMARC aggregate report to ElasticSearch
 
     Args:
         aggregate_report (OrderedDict): A parsed forensic report
         index_suffix (str): The suffix of the name of the index to save to
+        monthly_indexes (bool): Use monthly indexes instead of daily indexes
 
     Raises:
             AlreadySaved
@@ -286,7 +288,10 @@ def save_aggregate_report_to_elasticsearch(aggregate_report,
     end_date = human_timestamp_to_datetime(metadata["end_date"])
     begin_date_human = begin_date.strftime("%Y-%m-%d %H:%M:%S")
     end_date_human = end_date.strftime("%Y-%m-%d %H:%M:%S")
-    index_date = begin_date.strftime("%Y-%m-%d")
+    if monthly_indexes:
+        index_date = begin_date.strftime("%Y-%m")
+    else:
+        index_date = begin_date.strftime("%Y-%m-%d")
     aggregate_report["begin_date"] = begin_date
     aggregate_report["end_date"] = end_date
     date_range = [aggregate_report["begin_date"],
@@ -375,13 +380,16 @@ def save_aggregate_report_to_elasticsearch(aggregate_report,
 
 
 def save_forensic_report_to_elasticsearch(forensic_report,
-                                          index_suffix=None):
+                                          index_suffix=None,
+                                          monthly_indexes=False):
     """
         Saves a parsed DMARC forensic report to ElasticSearch
 
         Args:
             forensic_report (OrderedDict): A parsed forensic report
             index_suffix (str): The suffix of the name of the index to save to
+            monthly_indexes (bool): Use monthly indexes instead of daily
+                                    indexes
 
         Raises:
             AlreadySaved
@@ -485,7 +493,10 @@ def save_forensic_report_to_elasticsearch(forensic_report,
     index = "dmarc_forensic"
     if index_suffix:
         index = "{0}_{1}".format(index, index_suffix)
-    index_date = arrival_date.strftime("%Y-%m-%d")
+    if monthly_indexes:
+        index_date = arrival_date.strftime("%Y-%m")
+    else:
+        index_date = arrival_date.strftime("%Y-%m-%d")
     index = "{0}-{1}".format(index, index_date)
     create_indexes([index])
     forensic_doc.meta.index = index