Merge pull request #223 from tom-henderson/s3

Allow saving reports to S3

README.rst

@@ -128,11 +128,15 @@ For example
    token = HECTokenGoesHere
    index = email
 
+   [s3]
+   bucket = my-bucket
+   path = parsedmarc
+
 The full set of configuration options are:
 
 - ``general``
-    - ``save_aggregate`` - bool: Save aggregate report data to the Elasticsearch and/or Splunk
-    - ``save_forensic`` - bool: Save forensic report data to the Elasticsearch and/or Splunk
+    - ``save_aggregate`` - bool: Save aggregate report data to Elasticsearch, Splunk and/or S3
+    - ``save_forensic`` - bool: Save forensic report data to Elasticsearch, Splunk and/or S3
     - ``strip_attachment_payloads`` - bool: Remove attachment payloads from results
     - ``output`` - str: Directory to place JSON and CSV files in
     - ``offline`` - bool: Do not use online queries for geolocation or DNS
@@ -191,6 +195,9 @@ The full set of configuration options are:
     - ``subject`` - str: The Subject header to use in the email (Default: parsedmarc report)
     - ``attachment`` - str: The ZIP attachment filenames
     - ``message`` - str: The email message (Default: Please see the attached parsedmarc report.)
+- ``s3``
+    - ``bucket`` - str: The S3 bucket name
+    - ``path`` - int: The path to upload reports to (Default: /)
 
 
 .. warning::

docs/example.ini

@@ -18,3 +18,7 @@ ssl = False
 url = https://splunkhec.example.com
 token = HECTokenGoesHere
 index = email
+
+[s3]
+bucket = my-bucket
+path = parsedmarc

docs/index.rst

@@ -132,11 +132,15 @@ For example
    token = HECTokenGoesHere
    index = email
 
+   [s3]
+   bucket = my-bucket
+   path = parsedmarc
+
 The full set of configuration options are:
 
 - ``general``
-    - ``save_aggregate`` - bool: Save aggregate report data to the Elasticsearch and/or Splunk
-    - ``save_forensic`` - bool: Save forensic report data to the Elasticsearch and/or Splunk
+    - ``save_aggregate`` - bool: Save aggregate report data to the Elasticsearch, Splunk and/or S3
+    - ``save_forensic`` - bool: Save forensic report data to the Elasticsearch, Splunk and/or S3
     - ``strip_attachment_payloads`` - bool: Remove attachment payloads from results
     - ``output`` - str: Directory to place JSON and CSV files in
     - ``offline`` - bool: Do not use online queries for geolocation or DNS
@@ -200,7 +204,9 @@ The full set of configuration options are:
     - ``subject`` - str: The Subject header to use in the email (Default: parsedmarc report)
     - ``attachment`` - str: The ZIP attachment filenames
     - ``message`` - str: The email message (Default: Please see the attached parsedmarc report.)
-
+- ``s3``
+    - ``bucket`` - str: The S3 bucket name
+    - ``path`` - int: The path to upload reports to (Default: /)
 
 .. warning::
 

parsedmarc/cli.py

@@ -19,7 +19,7 @@ from tqdm import tqdm
 from parsedmarc import get_dmarc_reports_from_inbox, watch_inbox, \
     parse_report_file, get_dmarc_reports_from_mbox, elastic, kafkaclient, \
     splunk, save_output, email_results, ParserError, __version__, \
-    InvalidDMARCReport
+    InvalidDMARCReport, s3
 from parsedmarc.utils import is_mbox
 
 logger = logging.getLogger("parsedmarc")
@@ -79,6 +79,14 @@ def _main():
                 )
             except Exception as error_:
                 logger.error("Kafka Error: {0}".format(error_.__str__()))
+        if opts.s3_bucket:
+            try:
+                s3_client = s3.S3Client(
+                    bucket_name=opts.s3_bucket,
+                    bucket_path=opts.s3_path,
+                )
+            except Exception as error_:
+                logger.error("S3 Error: {0}".format(error_.__str__()))
         if opts.save_aggregate:
             for report in reports_["aggregate_reports"]:
                 try:
@@ -104,6 +112,11 @@ def _main():
                 except Exception as error_:
                     logger.error("Kafka Error: {0}".format(
                          error_.__str__()))
+                try:
+                    if opts.s3_bucket:
+                        s3_client.save_aggregate_report_to_s3(report)
+                except Exception as error_:
+                    logger.error("S3 Error: {0}".format(error_.__str__()))
             if opts.hec:
                 try:
                     aggregate_reports_ = reports_["aggregate_reports"]
@@ -138,6 +151,11 @@ def _main():
                 except Exception as error_:
                     logger.error("Kafka Error: {0}".format(
                         error_.__str__()))
+                try:
+                    if opts.s3_bucket:
+                        s3_client.save_forensic_report_to_s3(report)
+                except Exception as error_:
+                    logger.error("S3 Error: {0}".format(error_.__str__()))
             if opts.hec:
                 try:
                     forensic_reports_ = reports_["forensic_reports"]
@@ -241,6 +259,8 @@ def _main():
                      smtp_to=[],
                      smtp_subject="parsedmarc report",
                      smtp_message="Please see the attached DMARC results.",
+                     s3_bucket=None,
+                     s3_path=None,
                      log_file=args.log_file,
                      n_procs=1,
                      chunk_size=1
@@ -469,6 +489,22 @@ def _main():
                 opts.smtp_attachment = smtp_config["attachment"]
             if "message" in smtp_config:
                 opts.smtp_message = smtp_config["message"]
+        if "s3" in config.sections():
+            s3_config = config["s3"]
+            if "bucket" in s3_config:
+                opts.s3_bucket = s3_config["bucket"]
+            else:
+                logger.critical("bucket setting missing from the "
+                                "s3 config section")
+                exit(-1)
+            if "path" in s3_config:
+                opts.s3_path = s3_config["path"]
+                if opts.s3_path.startswith("/"):
+                    opts.s3_path = opts.s3_path[1:]
+                if opts.s3_path.endswith("/"):
+                    opts.s3_path = opts.s3_path[:-1]
+            else:
+                opts.s3_path = ""
 
     logging.basicConfig(level=logging.WARNING)
     logger.setLevel(logging.WARNING)

parsedmarc/s3.py

@@ -0,0 +1,68 @@
+# -*- coding: utf-8 -*-
+
+import logging
+import json
+import boto3
+
+from parsedmarc.utils import human_timestamp_to_datetime
+
+logger = logging.getLogger("parsedmarc")
+
+
+class S3Client(object):
+    """A client for a Amazon S3"""
+
+    def __init__(self, bucket_name, bucket_path):
+        """
+        Initializes the S3Client
+        Args:
+            bucket_name (str): The S3 Bucket
+            bucket_path (str): The path to save reports
+        """
+        self.bucket_name = bucket_name
+        self.bucket_path = bucket_path
+        self.metadata_keys = [
+            "org_name",
+            "org_email",
+            "report_id",
+            "begin_date",
+            "end_date",
+        ]
+
+        self.s3 = boto3.resource('s3')
+        self.bucket = self.s3.Bucket(self.bucket_name)
+
+    def save_aggregate_report_to_s3(self, report):
+        self.save_report_to_s3(report, 'aggregate')
+
+    def save_forensic_report_to_s3(self, report):
+        self.save_report_to_s3(report, 'forensic')
+
+    def save_report_to_s3(self, report, report_type):
+        report_date = human_timestamp_to_datetime(
+            report["report_metadata"]["begin_date"]
+        )
+        report_id = report["report_metadata"]["report_id"]
+        path_template = "{0}/{1}/year={2}/month={3:02d}/day={4:02d}/{5}.json"
+        object_path = path_template.format(
+            self.bucket_path,
+            report_type,
+            report_date.year,
+            report_date.month,
+            report_date.day,
+            report_id
+        )
+        logger.debug("Saving {0} report to s3://{1}/{2}".format(
+            report_type,
+            self.bucket_name,
+            object_path))
+        object_metadata = {
+            k: v
+            for k, v in report["report_metadata"].items()
+            if k in self.metadata_keys
+        }
+        self.bucket.put_object(
+            Body=json.dumps(report),
+            Key=object_path,
+            Metadata=object_metadata
+        )

requirements.txt

@@ -27,3 +27,4 @@ sphinx_rtd_theme>=0.4.3
 wheel>=0.33.6
 codecov>=2.0.15
 lxml>=4.4.0
+boto3>=1.16.63

setup.py

@@ -98,7 +98,8 @@ setup(
                       'elasticsearch-dsl>=7.2.0,<8.0.0',
                       'kafka-python>=1.4.4',
                       'tqdm>=4.31.1',
-                      'lxml>=4.4.0'
+                      'lxml>=4.4.0',
+                      'boto3>=1.16.63'
                       ],
 
     entry_points={