From 63dc7509a66d6c0653bae12a7ff311b0e8adfe71 Mon Sep 17 00:00:00 2001 From: Neil Smith Date: Mon, 4 Apr 2016 18:53:02 +0100 Subject: [PATCH] Added argment parsing, changed name of config file --- .gitignore | 2 +- dmarc_to_database | 21 +++++++++++++++++++-- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/.gitignore b/.gitignore index e763a5e..c4564ff 100644 --- a/.gitignore +++ b/.gitignore @@ -45,7 +45,7 @@ nosetests.xml *.log # Config file -dmarc.ini +*.ini # DMARC reports *xml diff --git a/dmarc_to_database b/dmarc_to_database index ac7baa3..ca6eafb 100755 --- a/dmarc_to_database +++ b/dmarc_to_database @@ -9,6 +9,16 @@ import xml.etree.ElementTree import psycopg2 import re import datetime +import argparse + +parser = argparse.ArgumentParser(description='Process DMARC records.') +parser.add_argument('-c', '--config', action='store', + default='', dest='config_file', + help='Path to config file') +parser.add_argument('-t', '--test', action='store_true', + default=False, + help='Test, but do not add records to the database') +args = parser.parse_args() def fetch_msg(num): return mailbox.uid('FETCH', num, '(RFC822)')[1][0][1] @@ -158,7 +168,13 @@ def write_report(connection, cursor, report): connection.commit() config = configparser.ConfigParser() -config.read('dmarc.ini') +if args.config_file: + config.read(args.config_file) +else: + config.read(['/etc/dmarc_to_database.ini', './dmarc_to_database.ini']) + +if not config.sections(): + raise RuntimeError('Could not find configuration file') conn = psycopg2.connect(host=config['database']['server'], database=config['database']['database'], @@ -196,6 +212,7 @@ for report in dmarc_reports: results = cur.fetchall() if not results: print('write', report.find('./report_metadata/report_id').text) - write_report(conn, cur, report) + if not args.test: + write_report(conn, cur, report) conn.close() -- 2.34.1