Added argment parsing, changed name of config file

diff --git a/.gitignore b/.gitignore
index e763a5e2a8fea6dd2e079930ac0a1a5c29326cf7..c4564ff210a36760c204a31cc12224852ad42d1b 100644 (file)
--- 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 ac7baa39c034ac7e00dc3616de4ac0c5be382c0f..ca6eafb52c20477651a6078a97f824cc8be86dc1 100755 (executable)
--- 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()