Add function to temporarily mute errors from an account

4 files changed, 29 insertions, 9 deletions

diff --git a/bin/emailcanary b/bin/emailcanary
index 706f633..46b47e9 100755
--- a/bin/emailcanary
+++ b/bin/emailcanary
@@ -3,6 +3,7 @@
 import argparse
 import sys
 import smtplib
+import time
 import traceback
 from emailcanary import canarydb
 from emailcanary import canary
@@ -20,6 +21,7 @@ def parse_args():
     group.add_argument('--chirp', metavar='listAddress', help='Send an email to the given canary list address.')
     group.add_argument('--check', metavar='listAddress', help='Check the recepient addresses for the list address.')
     group.add_argument('--add', nargs=4, metavar=('listAddress', 'address', 'imapserver', 'password'), help='Add a new recepient and list.')
+    group.add_argument('--mute', nargs=3, metavar=('listAddress', 'address', 'days'), help='Mute errors for the account for the given number of days.')
     group.add_argument('--remove', nargs='+', metavar=('listAddress', 'address'), help='Remove recepients from the list.')
     group.add_argument('--list', action='store_true', help='List all configured addresses.')
 
@@ -50,6 +52,10 @@ def add(db, listAddress, recepient, imapserver, password):
     db.add_account(listAddress, recepient, imapserver, password)
     return SUCCESS
 
+def mute(db, listAddress, recepient, days):
+    db.mute_account(listAddress, recepient, time.time() + 86400 * days)
+    return SUCCESS
+
 def remove(db, listAddress, recepients):
     if len(recepients) == 0:
         recepients = db.get_accounts(listAddress)
@@ -83,6 +89,8 @@ def main():
             return list(db)
         elif args.add:
             return add(db, args.add[0], args.add[1], args.add[2], args.add[3])
+        elif args.mute:
+            return mute(db, args.mute[0], args.mute[1], int(args.mute[2]))
         elif args.remove:
             return remove(db, args.remove[0], args.remove[1:])
         else:
diff --git a/emailcanary/canary.py b/emailcanary/canary.py
index 93e9d42..2be8ccf 100644
--- a/emailcanary/canary.py
+++ b/emailcanary/canary.py
@@ -26,7 +26,8 @@ class Canary:
         accounts = self.db.get_accounts(listAddress)
         if len(accounts) == 0:
             raise Exception("No receipients for listAddress '%s'", (listAddress,))
-        for (listAddress, address, imapserver, password) in accounts:
+        result = []
+        for (listAddress, address, imapserver, password, mute) in accounts:
             mail = emailutils.get_imap(imapserver, address, password)
             these_subjects = []
             for uid in emailutils.get_mail_uids(mail):
@@ -34,7 +35,9 @@ class Canary:
                 if message is not None and self.processMessage(address, message):
                     emailutils.delete_message(mail, uid)
             emailutils.close(mail)
-        return self.db.get_missing_pongs(listAddress)
+            if time.time() > mute:
+                result.extend(self.db.get_missing_pongs(listAddress, address))
+        return result
 
     def processMessage(self, receipient, msg):
         match = re.match('.*Canary Email (.+)', msg['Subject'])
diff --git a/emailcanary/canarydb.py b/emailcanary/canarydb.py
index 8bf944f..400fc6d 100644
--- a/emailcanary/canarydb.py
+++ b/emailcanary/canarydb.py
@@ -8,14 +8,14 @@ class CanaryDB:
         # Create Tables if necessary
         cursor = self.conn.cursor()
         cursor.execute("CREATE TABLE IF NOT EXISTS pings (listAddress, address, uuid, timesent timestamp, timereceived timestamp)")
-        cursor.execute("CREATE TABLE IF NOT EXISTS accounts (list, address, imapserver, password)")
+        cursor.execute("CREATE TABLE IF NOT EXISTS accounts (list, address, imapserver, password, mute)")
 
     def close(self):
         self.conn.close()
 
     def add_account(self, listAddress, address, imapserver, password):
         cursor = self.conn.cursor()
-        cursor.execute("INSERT INTO accounts (list, address, imapserver, password) VALUES (?, ?, ?, ?)", \
+        cursor.execute("INSERT INTO accounts (list, address, imapserver, password, mute) VALUES (?, ?, ?, ?, 0)", \
             (listAddress, address, imapserver, password))
         self.conn.commit()
 
@@ -24,19 +24,25 @@ class CanaryDB:
         cursor.execute("DELETE FROM accounts WHERE list=? AND address=?", (listAddress, address))
         self.conn.commit()
 
+    def mute_account(self, listAddress, address, until):
+        cursor = self.conn.cursor()
+        cursor.execute("UPDATE accounts SET mute=? WHERE list=? AND address=?", (until, listAddress, address))
+        self.conn.commit()
+
     def get_accounts(self, listAddress = None):
         cursor = self.conn.cursor()
         if listAddress:
-            cursor.execute("SELECT list, address, imapserver, password FROM accounts WHERE list=?", (listAddress,));
+            cursor.execute("SELECT list, address, imapserver, password, mute FROM accounts WHERE list=?", (listAddress,));
         else:
-            cursor.execute("SELECT list, address, imapserver, password FROM accounts");
+            cursor.execute("SELECT list, address, imapserver, password, mute FROM accounts");
         results = list()
         for row in cursor:
             listAddress = row[0]
             address = row[1]
             imapserver = row[2]
             password = row[3]
-            results.append((listAddress, address, imapserver, password))
+            mute = row[4]
+            results.append((listAddress, address, imapserver, password, mute))
         return results
 
     def get_recipients_for_list(self, listAddress):
@@ -59,11 +65,13 @@ class CanaryDB:
             (time, address, uuid))
         self.conn.commit()
 
-    def get_missing_pongs(self, listAddress = None):
+    def get_missing_pongs(self, listAddress = None, address = None):
         '''Return a list of tupls of missing pongs.
            Each tupl contains (listAddress, uuid, address, time since send)'''
         cursor = self.conn.cursor()
-        if listAddress:
+        if listAddress and address:
+            cursor.execute("SELECT listAddress, uuid, address, timesent FROM pings WHERE timereceived IS NULL AND listAddress = ? AND address = ?", (listAddress, address,));
+        elif listAddress:
             cursor.execute("SELECT listAddress, uuid, address, timesent FROM pings WHERE timereceived IS NULL AND listAddress = ?", (listAddress,));
         else:
             cursor.execute("SELECT listAddress, uuid, address, timesent FROM pings WHERE timereceived IS NULL");
diff --git a/tests/test_canarydb.py b/tests/test_canarydb.py
index 47bdc52..751c2cf 100644
--- a/tests/test_canarydb.py
+++ b/tests/test_canarydb.py
@@ -89,6 +89,7 @@ class TestCanaryDB(unittest.TestCase):
         self.assertEqual(address, accounts[0][1])
         self.assertEqual(imapserver, accounts[0][2])
         self.assertEqual(password, accounts[0][3])
+        self.assertEqual(0, accounts[0][4])
 
         # Remove the account
         self.db.remove_account(listAddress, address)