Adding Canary Check and Accounts

Adding the canary check method and tests for the Canary class.
Adding accounts to CanaryDB.

8 files changed, 237 insertions, 26 deletions

diff --git a/Makefile b/Makefile
index e365f62..d2c655e 100644
--- a/Makefile
+++ b/Makefile
@@ -32,4 +32,4 @@ check:
 clean:
 	$(PYTHON) setup.py clean
 	rm -rf build/ MANIFEST dist build my_program.egg-info deb_dist
-	find . -name '*.pyc' -delete
\ No newline at end of file
+	find . -name '*.pyc' -delete
diff --git a/emailcanary/canary.py b/emailcanary/canary.py
index d07e2e9..49972db 100644
--- a/emailcanary/canary.py
+++ b/emailcanary/canary.py
@@ -1,5 +1,5 @@
 import uuid, datetime, time
-import email
+import email, emailutils
 import re
 
 class Canary:
@@ -8,26 +8,39 @@ class Canary:
 		self.smtp = smtp
 		self.fromaddress = fromaddress
 
-	def chirp(self, list, expectedreceipients):
-		uuid = uuid.uuid4()
+	def chirp(self, listAddress):
+		chirpUUID = str(uuid.uuid4())
 		now = datetime.datetime.now()
 
-		self.send(list, now, uuid)
-		for dest in expectedreceipients:
-			self.db.ping(dest, now, uuid)
-
-	def echo(self, receipient, msg):
-		uuid = re.match('Canary Email (.+)', msg['Subject']).group(1)
-		now = datetime.datetime.now()
-
-		self.db.pong(receipient, now, uuid)
-
-
-	def send(self, dest, date, uuid):
+		self.send(listAddress, now, chirpUUID)
+		for dest in self.db.get_recipients_for_list(listAddress):
+			self.db.ping(dest, now, chirpUUID)
+
+	def check(self, listAddress):
+		for (listAddress, address, imapserver, password) in self.db.get_accounts(listAddress):
+			mail = emailutils.get_imap(imapserver, address, password)
+			these_subjects = []
+			for uid in emailutils.get_mail_uids(mail):
+				message = emailutils.get_message(mail, uid)
+				if self.processMessage(address, message):
+					emailutils.delete_message(mail, uid)
+			emailutils.close(mail)
+
+	def processMessage(self, receipient, msg):
+		match = re.match('Canary Email (.+)', msg['Subject'])
+		if match:
+			chirpUUID = match.group(1)
+			now = datetime.datetime.now()
+			self.db.pong(receipient, now, chirpUUID)
+			return True
+		return False
+
+
+	def send(self, dest, date, chirpUUID):
 		msg = email.message.Message()
 		msg['From'] = self.fromaddress
 		msg['To'] = dest
-		msg['Subject'] = "Canary Email " + str(uuid)
+		msg['Subject'] = "Canary Email " + chirpUUID
 		msg['Date'] = email.utils.formatdate(time.mktime(date.timetuple()))
 
-		self.smtp.sendmail(self.fromaddress, dest, msg.as_string())
\ No newline at end of file
+		self.smtp.sendmail(self.fromaddress, dest, msg.as_string())
diff --git a/emailcanary/canarydb.py b/emailcanary/canarydb.py
index 112b8f5..59a2ea2 100644
--- a/emailcanary/canarydb.py
+++ b/emailcanary/canarydb.py
@@ -8,10 +8,45 @@ class CanaryDB:
 		# Create Tables if necessary
 		cursor = self.conn.cursor()
 		cursor.execute("CREATE TABLE IF NOT EXISTS pings (address, uuid, timesent timestamp, timereceived timestamp)")
+		cursor.execute("CREATE TABLE IF NOT EXISTS accounts (list, address, imapserver, password)")
 
 	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 (?, ?, ?, ?)", \
+			(listAddress, address, imapserver, password))
+		self.conn.commit()
+
+	def remove_account(self, address):
+		cursor = self.conn.cursor()
+		cursor.execute("DELETE FROM accounts WHERE address=?", (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,));
+		else:
+			cursor.execute("SELECT list, address, imapserver, password 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))
+		return results
+
+	def get_recipients_for_list(self, listAddress):
+		cursor = self.conn.cursor()
+		cursor.execute("SELECT address FROM accounts WHERE list=?", (listAddress,));
+		results = list()
+		for row in cursor:
+			results.append(row[0])
+		return results
+
 	def ping(self, address, time, uuid):
 		cursor = self.conn.cursor()
 		cursor.execute("INSERT INTO pings (address, timesent, uuid) VALUES (?, ?, ?)", \
@@ -29,7 +64,6 @@ class CanaryDB:
 		   Each tupl contains (uuid, address, time since send)'''
 		cursor = self.conn.cursor()
 		cursor.execute("SELECT uuid, address, timesent FROM pings WHERE timereceived IS NULL");
-
 		now = datetime.now()
 		results = list()
 		for row in cursor:
@@ -37,5 +71,4 @@ class CanaryDB:
 			address = row[1]
 			delta = now - row[2]
 			results.append((uuid, address, delta))
-
 		return results
diff --git a/emailcanary/email-digest-sender.py b/emailcanary/email-digest-sender.py
index f1583c0..a7ff79f 100644
--- a/emailcanary/email-digest-sender.py
+++ b/emailcanary/email-digest-sender.py
@@ -10,7 +10,7 @@ youve_got_mail = False
 all_subjects = {}
 
 for account in ACCOUNTS:
-	mail = emailutils.get_imap(account)
+	mail = emailutils.get_imap(account[0], account[1], account[2])
 	these_subjects = []
 	for uid in emailutils.get_mail_uids(mail):
 		message = emailutils.get_message(mail, uid)
diff --git a/emailcanary/emailutils.py b/emailcanary/emailutils.py
index 6646c1e..b87f637 100644
--- a/emailcanary/emailutils.py
+++ b/emailcanary/emailutils.py
@@ -1,9 +1,9 @@
 import imaplib, email
 
-def get_imap(account):
+def get_imap(server, username, password):
 	'''Connect and login via IMAP'''
-	mail = imaplib.IMAP4_SSL(account[0])
-	mail.login(account[1],  account[2])
+	mail = imaplib.IMAP4_SSL(server)
+	mail.login(username, password)
 	return mail
 
 def get_mail_uids(mail):
@@ -16,4 +16,12 @@ def get_message(mail, uid):
 	'''Get a single email message object by UID'''
 	result, data = mail.uid('fetch', uid, '(RFC822)')
 	raw_email = data[0][1]
-	return email.message_from_string(raw_email)
\ No newline at end of file
+	return email.message_from_string(raw_email)
+
+def delete_message(mail, uid):
+	result = mail.uid('store', uid, '+FLAGS', '\\Deleted')
+
+def close(mail):
+	mail.expunge()
+	mail.close()
+	mail.logout()
diff --git a/requirements.txt b/requirements.txt
index e51ea6d..c36aac9 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
 setuptools
 pylint
-unittest2
\ No newline at end of file
+unittest2
+mock
\ No newline at end of file
diff --git a/tests/test_canary.py b/tests/test_canary.py
new file mode 100644
index 0000000..ff104f6
--- /dev/null
+++ b/tests/test_canary.py
@@ -0,0 +1,103 @@
+import unittest
+import mock
+import tempfile, shutil
+import datetime
+import smtplib, email
+
+from emailcanary import canary
+from emailcanary import canarydb
+
+FROM_ADDRESS = "from@example.com"
+LIST_ADDRESS = "list@example.com"
+
+USER_ADDRESS1 = "user1@example.com"
+USER_ADDRESS2 = "user2@example.com"
+
+SERVER = "mail.example.com"
+PASSWORD = "secret"
+
+class TestCanary(unittest.TestCase):
+	def setUp(self):
+		self.db = mock.Mock(canarydb.CanaryDB)
+		self.smtp = mock.Mock(smtplib.SMTP_SSL)
+		self.canary = canary.Canary(self.db, self.smtp, FROM_ADDRESS)
+		canary.emailutils.get_imap = mock.Mock()
+		canary.emailutils.get_message = mock.Mock()
+		canary.emailutils.get_mail_uids = mock.Mock()
+		canary.emailutils.delete_message = mock.Mock()
+		canary.emailutils.close = mock.Mock()
+
+	def tearDown(self):
+		pass
+
+	def testChirp(self):
+		# Setup Mock
+		self.db.get_recipients_for_list.return_value = [USER_ADDRESS1, USER_ADDRESS2]
+
+		# Test chirp
+		self.canary.chirp(LIST_ADDRESS)
+
+		# Assert DB updated
+		self.db.get_recipients_for_list.assert_called_with(LIST_ADDRESS)
+		self.db.ping.assert_has_calls( \
+			[mock.call(USER_ADDRESS1, mock.ANY, mock.ANY), \
+			 mock.call(USER_ADDRESS2, mock.ANY, mock.ANY)])
+		args = self.db.ping.call_args
+		expectedSubject = "Canary Email " + args[0][2]
+
+		# Assert emails were sent
+		self.assertEqual(1, self.smtp.sendmail.call_count)
+		args = self.smtp.sendmail.call_args[0]
+		self.assertEqual(FROM_ADDRESS, args[0])
+		self.assertEqual(LIST_ADDRESS, args[1])
+		msg = email.message_from_string(args[2])
+		self.assertEqual(FROM_ADDRESS, msg['From'])
+		self.assertEqual(LIST_ADDRESS, msg['To'])
+		self.assertEqual(expectedSubject, msg['Subject'])
+
+	def testCheck(self):
+		# Setup mocks
+		expectedUUID = "1234-5678-9012-3456"
+		self.db.get_accounts.return_value = [ \
+			(LIST_ADDRESS, USER_ADDRESS1, SERVER, PASSWORD), \
+			(LIST_ADDRESS, USER_ADDRESS2, SERVER, PASSWORD)]
+		canary.emailutils.get_mail_uids.return_value = [1]
+		canary.emailutils.get_message.return_value = {'Subject': "Canary Email " + expectedUUID}
+
+		# Test check
+		self.canary.check(LIST_ADDRESS)
+
+		# Assert DB calls
+		self.db.get_accounts.assert_called_with(LIST_ADDRESS)
+		self.db.pong.assert_has_calls([ \
+			mock.call(USER_ADDRESS1, mock.ANY, expectedUUID), \
+			mock.call(USER_ADDRESS2, mock.ANY, expectedUUID)])
+
+		# Assert mail calls
+		canary.emailutils.get_imap.assert_has_calls([ \
+			mock.call(SERVER, USER_ADDRESS1, PASSWORD), \
+			mock.call(SERVER, USER_ADDRESS2, PASSWORD)])
+		canary.emailutils.get_message.assert_called_with(canary.emailutils.get_imap.return_value, 1)
+		canary.emailutils.delete_message.assert_called_with(canary.emailutils.get_imap.return_value, 1)
+		canary.emailutils.close.assert_called_with(canary.emailutils.get_imap.return_value)
+
+	def testDontDeleteOtherMail(self):
+		# Setup mocks
+		self.db.get_accounts.return_value = [(LIST_ADDRESS, USER_ADDRESS1, SERVER, PASSWORD)]
+		canary.emailutils.get_mail_uids.return_value = [1]
+		canary.emailutils.get_message.return_value = {'Subject': "Buy Our New Widget"}
+
+		# Test check
+		self.canary.check(LIST_ADDRESS)
+
+		# Assert DB calls
+		self.db.get_accounts.assert_called_with(LIST_ADDRESS)
+		self.db.pong.assert_not_called()
+
+		# Assert mail calls
+		canary.emailutils.get_imap.assert_called_with(SERVER, USER_ADDRESS1, PASSWORD)
+		canary.emailutils.get_message.assert_called_with(canary.emailutils.get_imap.return_value, 1)
+		canary.emailutils.delete_message.assert_not_called()
+		canary.emailutils.close.assert_called_with(canary.emailutils.get_imap.return_value)
+
+
diff --git a/tests/test_canarydb.py b/tests/test_canarydb.py
index 5a0fc23..8a29d23 100644
--- a/tests/test_canarydb.py
+++ b/tests/test_canarydb.py
@@ -64,3 +64,56 @@ class TestCanaryDB(unittest.TestCase):
 		self.assertEqual(address, firstMissing[1])
 		delta = firstMissing[2].total_seconds() - expectedDelta.total_seconds()
 		self.assertTrue(delta <= 10)
+
+	def testAccounts(self):
+		listAddress = "list@example.org"
+		address = "user@example.net"
+		imapserver = "imap.example.net"
+		password = "secretpassword"
+
+		# Verify that no accounts exist
+		accounts = self.db.get_accounts()
+		self.assertEqual(0, len(accounts))
+
+		# Add one account
+		self.db.add_account(listAddress, address, imapserver, password)
+
+		# Verify that the account exists
+		accounts = self.db.get_accounts()
+		self.assertEqual(1, len(accounts))
+		self.assertEqual(listAddress, accounts[0][0])
+		self.assertEqual(address, accounts[0][1])
+		self.assertEqual(imapserver, accounts[0][2])
+		self.assertEqual(password, accounts[0][3])
+
+		# Remove the account
+		self.db.remove_account(address)
+		accounts = self.db.get_accounts()
+		self.assertEqual(0, len(accounts))
+
+	def testGetRecipientsForList(self):
+		listAddress1 = "list1@example.org"
+		listAddress2 = "list2@example.org"
+		imapserver = "imap.example.net"
+		password = "secretpassword"
+		address1 = "user1@example.net"
+		address2 = "user2@example.net"
+
+		# No accounts
+		self.assertEqual([], self.db.get_recipients_for_list(listAddress1));
+		self.assertEqual([], self.db.get_recipients_for_list(listAddress2));
+
+		# One account
+		self.db.add_account(listAddress1, address1, imapserver, password)
+		self.assertEqual([address1], self.db.get_recipients_for_list(listAddress1));
+		self.assertEqual([], self.db.get_recipients_for_list(listAddress2));
+
+		# Two accounts
+		self.db.add_account(listAddress1, address2, imapserver, password)
+		self.assertEqual([address1, address2], self.db.get_recipients_for_list(listAddress1));
+		self.assertEqual([], self.db.get_recipients_for_list(listAddress2));
+
+		# Two lists
+		self.db.add_account(listAddress2, address1, imapserver, password)
+		self.assertEqual([address1, address2], self.db.get_recipients_for_list(listAddress1));
+		self.assertEqual([address1], self.db.get_recipients_for_list(listAddress2));