X-Git-Url: https://git.njae.me.uk/?a=blobdiff_plain;f=dmarc_to_database.py;h=46249649141088ccf310859496d7f36ac9a41326;hb=fa101ebf71cd46c51df29eac35a001f60dd6434e;hp=2fd237568b03a466ac5ed2f16d56215344d8ce23;hpb=4501f68fa04a34934ff3c3ec8e9ddf43e060c5e5;p=dmarc.git

diff --git a/dmarc_to_database.py b/dmarc_to_database.py
index 2fd2375..4624964 100644
--- a/dmarc_to_database.py
+++ b/dmarc_to_database.py
@@ -1,3 +1,5 @@
+#!/usr/bin/python3
+
 import configparser
 import imaplib
 import email
@@ -155,13 +157,14 @@ def write_report(connection, cursor, report):
 config = configparser.ConfigParser()
 config.read('dmarc.ini')
 
-with psycopg2.connect(host=config['database']['server'],
+conn = psycopg2.connect(host=config['database']['server'],
                         database=config['database']['database'], 
                         user=config['database']['username'], 
-                        password=config['database']['password']) as conn:
-    with conn.cursor() as cur:
-        cur.execute('select max(report_metadata_date_range_end) from reports')
-        results = cur.fetchall()
+                        password=config['database']['password']) 
+
+cur = conn.cursor()
+cur.execute('select max(report_metadata_date_range_end) from reports')
+results = cur.fetchall()
 most_recent_date = results[0][0]
 
 mailbox = imaplib.IMAP4(host=config['imap']['server'], 
@@ -184,19 +187,12 @@ dmarc_reports = [report for report_set in [extract_report(fetch_msg(n)) for n in
 mailbox.close()
 mailbox.logout()
 
+for report in dmarc_reports:
+    cur.execute('select id, report_metadata_report_id from reports where report_metadata_report_id = %s;', 
+        [report.find('./report_metadata/report_id').text])
+    results = cur.fetchall()
+    if not results:
+        print('write', report.find('./report_metadata/report_id').text)
+        write_report(conn, cur, report)
 
-with psycopg2.connect(host=config['database']['server'],
-                        database=config['database']['database'], 
-                        user=config['database']['username'], 
-                        password=config['database']['password']) as conn:
-    with conn.cursor() as cur:
-        for report in dmarc_reports:
-            cur.execute('select id, report_metadata_report_id from reports where report_metadata_report_id = %s;', 
-                [report.find('./report_metadata/report_id').text])
-            results = cur.fetchall()
-            if not results:
-                print('write', report.find('./report_metadata/report_id').text)
-                write_report(conn, cur, report)
-
-
-
+conn.close()