Merge pull request #21 from mikesiegel/mikesiegel_kafka

Add Kafka Support

parsedmarc/__init__.py

@@ -44,7 +44,7 @@ import imapclient.exceptions
 import dateparser
 import mailparser
 
-__version__ = "4.2.0"
+__version__ = "4.2.0k"
 
 logger = logging.getLogger("parsedmarc")
 

parsedmarc/cli.py

@@ -13,20 +13,26 @@ import json
 from elasticsearch.exceptions import ElasticsearchException
 
 from parsedmarc import IMAPError, get_dmarc_reports_from_inbox, \
-    parse_report_file, elastic, splunk, save_output, watch_inbox, \
-    email_results, SMTPError, ParserError, __version__
+    parse_report_file, elastic, kafkaclient, splunk, save_output, \
+    watch_inbox, email_results, SMTPError, ParserError, __version__
 
 logger = logging.getLogger("parsedmarc")
 
 
 def _main():
     """Called when the module is executed"""
+
     def process_reports(reports_):
         output_str = "{0}\n".format(json.dumps(reports_,
                                                ensure_ascii=False,
                                                indent=2))
         if not args.silent:
             print(output_str)
+        if args.kafka_hosts:
+            try:
+                kafkaClient = kafkaclient.KafkaClient(args.kafka_hosts)
+            except Exception as error:
+                logger.error("Kafka Error: {0}".format(error.__str__()))
         if args.save_aggregate:
             for report in reports_["aggregate_reports"]:
                 try:
@@ -39,6 +45,13 @@ def _main():
                     logger.error("Elasticsearch Error: {0}".format(
                         error_.__str__()))
                     exit(1)
+                try:
+                    if args.kafka_hosts:
+                        kafkaClient.save_aggregate_reports_to_kafka(
+                            report, kafka_aggregate_topic)
+                except Exception as error_:
+                    logger.error("Kafka Error: {0}".format(
+                         error_.__str__()))
             if args.hec:
                 try:
                     aggregate_reports_ = reports_["aggregate_reports"]
@@ -58,6 +71,14 @@ def _main():
                 except ElasticsearchException as error_:
                     logger.error("Elasticsearch Error: {0}".format(
                         error_.__str__()))
+                try:
+                    if args.kafka_hosts:
+                        kafkaClient.save_forensic_reports_to_kafka(
+                            report, kafka_forensic_topic)
+
+                except Exception as error_:
+                    logger.error("Kafka Error: {0}".format(
+                        error_.__str__()))
             if args.hec:
                 try:
                     forensic_reports_ = reports_["forensic_reports"]
@@ -126,6 +147,15 @@ def _main():
                             default=False,
                             help="Skip certificate verification for Splunk "
                                  "HEC")
+    arg_parser.add_argument("-K", "--kafka-hosts", nargs="*",
+                            help="A list of one or more Kafka hostnames"
+                            " or URLs")
+    arg_parser.add_argument("--kafka-aggregate-topic",
+                            help="The Kafka topic to publish aggregate "
+                            "reports to.")
+    arg_parser.add_argument("--kafka-forensic_topic",
+                            help="The Kafka topic to publish forensic reports"
+                            " to.")
     arg_parser.add_argument("--save-aggregate", action="store_true",
                             default=False,
                             help="Save aggregate reports to search indexes")
@@ -196,7 +226,8 @@ def _main():
         es_forensic_index = "{0}_{1}".format(es_forensic_index, suffix)
 
     if args.save_aggregate or args.save_forensic:
-        if args.elasticsearch_host is None and args.hec is None:
+        if (args.elasticsearch_host is None and args.hec
+                and args.kafka_hosts is None):
             args.elasticsearch_host = ["localhost:9200"]
         try:
             if args.elasticsearch_host:
@@ -219,6 +250,15 @@ def _main():
                                       args.hec_index,
                                       verify=verify)
 
+    kafka_aggregate_topic = "dmarc_aggrregate"
+    kafka_forensic_topic = "dmarc_forensic"
+
+    if args.kafka_aggregate_topic:
+        kafka_aggregate_topic = args.kafka_aggregate_topic
+
+    if args.kafka_forensic_topic:
+        kafka_forensic_topic = args.kafka_forensic_topic
+
     file_paths = []
     for file_path in args.file_path:
         file_paths += glob(file_path)

parsedmarc/kafkaclient.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+from kafka import KafkaProducer
+from kafka.errors import NoBrokersAvailable, UnknownTopicOrPartitionError
+import json
+
+
+class KafkaError(RuntimeError):
+    """Raised when a Kafka error occurs"""
+
+
+class KafkaClient(object):
+    def __init__(self, kafka_hosts):
+        try:
+            def serializer(v): lambda v: json.dumps(v).encode('utf-8')
+            self.producer = KafkaProducer(
+                          value_serializer=serializer,
+                          bootstrap_servers=kafka_hosts)
+        except NoBrokersAvailable:
+            raise KafkaError("No Kafka brokers availabe")
+
+    def save_aggregate_reports_to_kafka(self, aggregate_reports,
+                                        aggregate_topic):
+        """
+        Saves aggregate DMARC reports to Kafka
+
+        Args:
+            aggregate_reports (list):  A list of aggregate report dictionaries
+            to save to kafka
+
+        """
+        if type(aggregate_reports) == dict:
+            aggregate_reports = [aggregate_reports]
+
+        if len(aggregate_reports) < 1:
+            return
+
+        try:
+            self.producer.send(aggregate_topic, aggregate_reports)
+        except UnknownTopicOrPartitionError:
+                raise KafkaError("Unknown topic or partition on broker")
+        self.producer.flush()
+
+    def save_forensic_reports_to_kafka(self, forensic_reports, forensic_topic):
+            """
+            Saves forensic DMARC reports to Kafka
+
+            Args:
+                forensic_reports (list):  A list of forensic report dicts
+                to save to kafka
+
+            """
+            if type(forensic_reports) == dict:
+                forensic_reports = [forensic_reports]
+
+            if len(forensic_reports) < 1:
+                return
+
+            try:
+                self.producer.send(forensic_topic, forensic_reports)
+            except UnknownTopicOrPartitionError:
+                raise KafkaError("Unknown topic or partition on broker")
+            self.producer.flush()

requirements.txt

@@ -16,3 +16,4 @@ sphinx_rtd_theme
 collective.checkdocs
 wheel
 rstcheck
+kafka-python

setup.py

@@ -14,7 +14,7 @@ from setuptools import setup
 from codecs import open
 from os import path
 
-__version__ = "4.2.0"
+__version__ = "4.2.0k"
 description = "A Python package and CLI for parsing aggregate and " \
               "forensic DMARC reports"