Tidied files

[dmarc.git] / dmarc_to_database

diff --git a/dmarc_to_database b/dmarc_to_database
index a8141fdc108478097cc290e86890901c59cb6fd5..38ec02837ae4c6d4caf8ffb080c755ebd5c4a4dd 100755 (executable)
--- a/dmarc_to_database
+++ b/dmarc_to_database
@@ -9,15 +9,28 @@ 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]
 
 def xml_of_part(part):
-    with zipfile.ZipFile(io.BytesIO(part.get_payload(decode=True))) as zf:
+    try:
+        with zipfile.ZipFile(io.BytesIO(part.get_payload(decode=True))) as zf:
             fn = zf.infolist()[0].filename
             contents = zf.read(fn).decode('utf-8')
             return xml.etree.ElementTree.fromstring(contents)
+    except zipfile.BadZipFile:
+        return None
 
 
 def xml_of(message):
@@ -87,10 +100,10 @@ field_maps = {'./policy_published/adkim': {'pg_field_name': 'policy_published_ad
   'pg_type': 'inet'},
  './report_metadata/date_range/begin': {'pg_field_name': 'report_metadata_date_range_begin',
   'pg_table': 'reports',
-  'pg_type': 'timestamp'},
+  'pg_type': 'timestamptz'},
  './report_metadata/date_range/end': {'pg_field_name': 'report_metadata_date_range_end',
   'pg_table': 'reports',
-  'pg_type': 'timestamp'},
+  'pg_type': 'timestamptz'},
  './report_metadata/email': {'pg_field_name': 'report_metadata_email',
   'pg_table': 'reports',
   'pg_type': 'varchar'},
@@ -117,7 +130,7 @@ def build_insert_command(table_name, report, preamble_values=None, i=None):
         field_names += [field_maps[f]['pg_field_name']]
         if field_maps[f]['pg_type'] == 'int':
             values[field_maps[f]['pg_field_name']] = int(report.find(fp).text)
-        elif field_maps[f]['pg_type'] == 'timestamp':
+        elif field_maps[f]['pg_type'] == 'timestamptz':
             # values[field_maps[f]['pg_field_name']] = datetime.datetime.utcfromtimestamp(int(report.find(fp).text))
             values[field_maps[f]['pg_field_name']] = \
                 datetime.datetime.fromtimestamp(int(report.find(fp).text),  
@@ -158,7 +171,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'], 
@@ -185,7 +204,8 @@ resp, nums = mailbox.uid('SEARCH', None, mails_from)
 
 
 dmarc_reports = [report for report_set in [extract_report(fetch_msg(n)) for n in nums[0].split()]
-                for report in report_set]
+                for report in report_set
+                if report]
 
 mailbox.close()
 mailbox.logout()
@@ -196,6 +216,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()