diff --git a/parsedmarc/cli.py b/parsedmarc/cli.py index 65a88e7a..c155e752 100644 --- a/parsedmarc/cli.py +++ b/parsedmarc/cli.py @@ -16,6 +16,9 @@ from parsedmarc import logger, IMAPError, get_dmarc_reports_from_inbox, \ parse_report_file, elastic, kafkaclient, splunk, save_output, watch_inbox, \ email_results, SMTPError, ParserError, __version__ +import sys +import os + def _main(): """Called when the module is executed""" def process_reports(reports_): diff --git a/parsedmarc/kafkaclient.py b/parsedmarc/kafkaclient.py index 55f6dfad..b1e27730 100644 --- a/parsedmarc/kafkaclient.py +++ b/parsedmarc/kafkaclient.py @@ -2,6 +2,7 @@ # -*- coding: utf-8 -*- from kafka import KafkaProducer +from collections import OrderedDict import json class KafkaError(RuntimeError): @@ -9,14 +10,12 @@ class KafkaError(RuntimeError): class KafkaClient(object): def __init__(self, kafka_hosts): - """ Right now lets just do one host""" - self.host = kafka_hosts - self.producer = KafkaProducer(value_serializer=lambda v: json.dumps(v).encode('utf-8'), + self.producer = KafkaProducer(value_serializer=lambda v: json.dumps(v).encode('utf-8'), bootstrap_servers=kafka_hosts) def save_aggregate_reports_to_kafka(self, aggregate_reports, aggregate_topic): """ - Saves aggregate DMARC reports to Splunk + Saves aggregate DMARC reports to Kafka Args: aggregate_reports (list): A list of aggregate report dictionaries @@ -29,9 +28,11 @@ class KafkaClient(object): if len(aggregate_reports) < 1: return - print("Report is {}".format(aggregate_reports)) - print("Report type is {}".format(type(aggregate_reports))) - self.producer.send(aggregate_topic, aggregate_reports) + for record in aggregate_reports['records']: + buffer = OrderedDict([('xml_schema', aggregate_reports['xml_schema']), + ('report_metadata',aggregate_reports['report_metadata']), + ('records', record)]) + self.producer.send(aggregate_topic, buffer) self.producer.flush()