Fixed error handling

[dmarc.git] / dmarc_to_database

diff --git a/dmarc_to_database b/dmarc_to_database
index ca6eafb52c20477651a6078a97f824cc8be86dc1..62d6a077995ba90482d13969a4aa0c08844d8e00 100755 (executable)
--- a/dmarc_to_database
+++ b/dmarc_to_database
@@ -24,10 +24,13 @@ 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):
@@ -148,24 +151,27 @@ def build_insert_command(table_name, report, preamble_values=None, i=None):
 
 
 def write_report(connection, cursor, report):
-    insert_string, values = build_insert_command('reports', report)
-    # print(insert_string, values)
-    cursor.execute(insert_string, values)
-    
-    for i in range(1, len(report.findall('./record'))+1):
-        field_names = []
-        cursor.execute('select id, report_metadata_report_id from reports where report_metadata_report_id = %s;', 
-            [report.find('./report_metadata/report_id').text])
-        results = cursor.fetchall()
-        if len(results) != 1:
-            raise RuntimeError('Could not find report record for report item')
-        else:
-            report_id = results[0][0]
-        insert_string, values = build_insert_command('report_items', report, i=i,
-                                                     preamble_values={'report_id': report_id})
+    try:
+        insert_string, values = build_insert_command('reports', report)
         # print(insert_string, values)
         cursor.execute(insert_string, values)
-    connection.commit()
+        
+        for i in range(1, len(report.findall('./record'))+1):
+            field_names = []
+            cursor.execute('select id, report_metadata_report_id from reports where report_metadata_report_id = %s;', 
+                [report.find('./report_metadata/report_id').text])
+            results = cursor.fetchall()
+            if len(results) != 1:
+                raise RuntimeError('Could not find report record for report item')
+            else:
+                report_id = results[0][0]
+            insert_string, values = build_insert_command('report_items', report, i=i,
+                                                         preamble_values={'report_id': report_id})
+            # print(insert_string, values)
+            cursor.execute(insert_string, values)
+        connection.commit()
+    except AttributeError:
+        pass
 
 config = configparser.ConfigParser()
 if args.config_file:
@@ -201,7 +207,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()