Merge pull request #315 from itkovian/auth-ssl-mail

stdweird · web-flow · commit e37f183dd85d · 2021-05-31T20:21:33.000+02:00
feat: allow authenticated + ssl mail
diff --git a/lib/vsc/utils/mail.py b/lib/vsc/utils/mail.py
@@ -34,12 +34,11 @@
 import logging
 import re
 import smtplib
-
+import ssl
 from email.mime.multipart import MIMEMultipart
 from email.mime.text import MIMEText
 from email.mime.image import MIMEImage
 
-
 class VscMailError(Exception):
     """Raised if the sending of an email fails for some reason."""
 
@@ -68,14 +67,43 @@ def __init__(self, mail_host=None, mail_to=None, mail_from=None, mail_subject=No
 class VscMail(object):
     """Class providing functionality to send out mail."""
 
-    def __init__(self, mail_host=None):
+    def __init__(
+        self,
+        mail_host='',
+        mail_port=0,
+        smtp_auth_user=None,
+        smtp_auth_password=None,
+        smtp_use_starttls=False):
+
         self.mail_host = mail_host
+        self.mail_port = mail_port
+        self.smtp_auth_user = smtp_auth_user
+        self.smtp_auth_password = smtp_auth_password
+        self.smtp_use_starttls = smtp_use_starttls
+
+    def _connect(self):
+        """
+        Connect to the mail host on the given port.
+
+        If provided, use authentication and TLS.
+        """
+        logging.debug("Using mail host %s, mail port %d", self.mail_host, self.mail_port)
+        s = smtplib.SMTP(host=self.mail_host, port=self.mail_port)
 
-    def _send(self,
-              mail_from,
-              mail_to,
-              mail_subject,
-              msg):
+        if self.smtp_use_starttls:
+            context = ssl.create_default_context()
+            s.starttls(context=context)
+            logging.debug("Started TLS connection")
+
+        if self.smtp_auth_user and self.smtp_auth_password:
+            s.login(user=self.smtp_auth_user, password=self.smtp_auth_password)
+            logging.debug("Authenticated")
+
+        s.connect()
+
+        return s
+
+    def _send(self, mail_from, mail_to, mail_subject, msg):
         """Actually send the mail.
 
         @type mail_from: string representing the sender.
@@ -85,13 +113,8 @@ def _send(self,
         """
 
         try:
-            if self.mail_host:
-                logging.debug("Using %s as the mail host", self.mail_host)
-                s = smtplib.SMTP(self.mail_host)
-            else:
-                logging.debug("Using the default mail host")
-                s = smtplib.SMTP()
-                s.connect()
+            s = self._connect()
+
             try:
                 s.sendmail(mail_from, mail_to, msg.as_string())
             except smtplib.SMTPHeloError as err:
@@ -105,27 +128,27 @@ def _send(self,
                 raise
             except smtplib.SMTPDataError as err:
                 raise
+
         except smtplib.SMTPConnectError as err:
             logging.exception("Cannot connect to the SMTP host %s", self.mail_host)
-            raise VscMailError(mail_host=self.mail_host,
-                               mail_to=mail_to,
-                               mail_from=mail_from,
-                               mail_subject=mail_subject,
-                               err=err)
+            raise VscMailError(
+                mail_host=self.mail_host,
+                mail_to=mail_to,
+                mail_from=mail_from,
+                mail_subject=mail_subject,
+                err=err)
         except Exception as err:
             logging.exception("Some unknown exception occurred in VscMail.sendTextMail. Raising a VscMailError.")
-            raise VscMailError(mail_host=self.mail_host,
-                               mail_to=mail_to,
-                               mail_from=mail_from,
-                               mail_subject=mail_subject,
-                               err=err)
-
-    def sendTextMail(self,
-                     mail_to,
-                     mail_from,
-                     reply_to,
-                     mail_subject,
-                     message):
+            raise VscMailError(
+                mail_host=self.mail_host,
+                mail_to=mail_to,
+                mail_from=mail_from,
+                mail_subject=mail_subject,
+                err=err)
+        else:
+            s.quit()
+
+    def sendTextMail(self, mail_to, mail_from, reply_to, mail_subject, message):
         """Send out the given message by mail to the given recipient(s).
 
         @type mail_to: string or list of strings
@@ -174,15 +197,16 @@ def _replace_images_cid(self, html, images):
 
         return html
 
-    def sendHTMLMail(self,
-                     mail_to,
-                     mail_from,
-                     reply_to,
-                     mail_subject,
-                     html_message,
-                     text_alternative,
-                     images=None,
-                     css=None):
+    def sendHTMLMail(
+        self,
+        mail_to,
+        mail_from,
+        reply_to,
+        mail_subject,
+        html_message,
+        text_alternative,
+        images=None,
+        css=None):
         """
         Send an HTML email message, encoded in a MIME/multipart message.
 
@@ -222,7 +246,7 @@ def sendHTMLMail(self,
 
         # Create the body of the message (a plain-text and an HTML version).
         if images is not None:
-            html_message = self.replace_images_cid(html_message, images)
+            html_message = self._replace_images_cid(html_message, images)
 
         # Record the MIME types of both parts - text/plain and text/html_message.
         msg_plain = MIMEText(text_alternative, 'plain')
diff --git a/test/exceptions.py b/test/exceptions.py
@@ -62,8 +62,9 @@ def test_loggedexception_defaultlogger(self):
         logToFile(tmplog, enable=False)
 
         log_re = re.compile("^%s :: BOOM( \(at .*:[0-9]+ in raise_loggedexception\))?$" % getRootLoggerName(), re.M)
-        logtxt = open(tmplog, 'r').read()
-        self.assertTrue(log_re.match(logtxt), "%s matches %s" % (log_re.pattern, logtxt))
+        with open(tmplog, 'r') as f:
+            logtxt = f.read()
+            self.assertTrue(log_re.match(logtxt), "%s matches %s" % (log_re.pattern, logtxt))
 
         # test formatting of message
         self.assertErrorRegex(LoggedException, 'BOOMBAF', raise_loggedexception, 'BOOM%s', 'BAF')
@@ -92,8 +93,9 @@ def test_loggedexception_specifiedlogger(self):
 
         rootlog = getRootLoggerName()
         log_re = re.compile("^%s.testlogger_one :: BOOM( \(at .*:[0-9]+ in raise_loggedexception\))?$" % rootlog, re.M)
-        logtxt = open(tmplog, 'r').read()
-        self.assertTrue(log_re.match(logtxt), "%s matches %s" % (log_re.pattern, logtxt))
+        with open(tmplog, 'r') as f:
+            logtxt = f.read()
+            self.assertTrue(log_re.match(logtxt), "%s matches %s" % (log_re.pattern, logtxt))
 
         os.remove(tmplog)
 
@@ -114,8 +116,9 @@ def test_loggedexception_callerlogger(self):
 
         rootlog = getRootLoggerName()
         log_re = re.compile("^%s(.testlogger_local)? :: BOOM( \(at .*:[0-9]+ in raise_loggedexception\))?$" % rootlog)
-        logtxt = open(tmplog, 'r').read()
-        self.assertTrue(log_re.match(logtxt), "%s matches %s" % (log_re.pattern, logtxt))
+        with open(tmplog, 'r') as f:
+            logtxt = f.read()
+            self.assertTrue(log_re.match(logtxt), "%s matches %s" % (log_re.pattern, logtxt))
 
         os.remove(tmplog)
 
@@ -143,12 +146,12 @@ def raise_testexception(msg, *args, **kwargs):
         rootlogname = getRootLoggerName()
 
         log_re = re.compile("^%s :: BOOM$" % rootlogname, re.M)
-        logtxt = open(tmplog, 'r').read()
-        self.assertTrue(log_re.match(logtxt), "%s matches %s" % (log_re.pattern, logtxt))
+        with open(tmplog, 'r') as f:
+            logtxt = f.read()
+            self.assertTrue(log_re.match(logtxt), "%s matches %s" % (log_re.pattern, logtxt))
 
-        f = open(tmplog, 'w')
-        f.write('')
-        f.close()
+        with open(tmplog, 'w') as f:
+            f.write('')
 
         # location is included if LOC_INFO_TOP_PKG_NAMES is defined
         TestException.LOC_INFO_TOP_PKG_NAMES = ['vsc']
@@ -157,12 +160,12 @@ def raise_testexception(msg, *args, **kwargs):
         logToFile(tmplog, enable=False)
 
         log_re = re.compile(r"^%s :: BOOM \(at (?:.*?/)?vsc/install/testing.py:[0-9]+ in assertErrorRegex\)$" % rootlogname)
-        logtxt = open(tmplog, 'r').read()
-        self.assertTrue(log_re.match(logtxt), "%s matches %s" % (log_re.pattern, logtxt))
+        with open(tmplog, 'r') as f:
+            logtxt = f.read()
+            self.assertTrue(log_re.match(logtxt), "%s matches %s" % (log_re.pattern, logtxt))
 
-        f = open(tmplog, 'w')
-        f.write('')
-        f.close()
+        with open(tmplog, 'w') as f:
+            f.write('')
 
         # absolute path of location is included if there's no match in LOC_INFO_TOP_PKG_NAMES
         TestException.LOC_INFO_TOP_PKG_NAMES = ['foobar']
@@ -171,8 +174,9 @@ def raise_testexception(msg, *args, **kwargs):
         logToFile(tmplog, enable=False)
 
         log_re = re.compile(r"^%s :: BOOM \(at (?:.*?/)?vsc/install/testing.py:[0-9]+ in assertErrorRegex\)$" % rootlogname)
-        logtxt = open(tmplog, 'r').read()
-        self.assertTrue(log_re.match(logtxt), "%s matches %s" % (log_re.pattern, logtxt))
+        with open(tmplog, 'r') as f:
+            logtxt = f.read()
+            self.assertTrue(log_re.match(logtxt), "%s matches %s" % (log_re.pattern, logtxt))
 
         os.remove(tmplog)
 
diff --git a/test/mail.py b/test/mail.py
@@ -0,0 +1,95 @@
+#
+# Copyright 2014-2021 Ghent University
+#
+# This file is part of vsc-base,
+# originally created by the HPC team of Ghent University (http://ugent.be/hpc/en),
+# with support of Ghent University (http://ugent.be/hpc),
+# the Flemish Supercomputer Centre (VSC) (https://www.vscentrum.be),
+# the Flemish Research Foundation (FWO) (http://www.fwo.be/en)
+# and the Department of Economy, Science and Innovation (EWI) (http://www.ewi-vlaanderen.be/en).
+#
+# https://github.com/hpcugent/vsc-base
+#
+# vsc-base is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Library General Public License as
+# published by the Free Software Foundation, either version 2 of
+# the License, or (at your option) any later version.
+#
+# vsc-base is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Library General Public License for more details.
+#
+# You should have received a copy of the GNU Library General Public License
+# along with vsc-base. If not, see <http://www.gnu.org/licenses/>.
+#
+"""
+Unit tests for the mail wrapper.
+
+@author: Andy Georges (Ghent University)
+"""
+import mock
+import os
+
+from vsc.install.testing import TestCase
+
+from email.mime.text import MIMEText
+from vsc.utils.mail import VscMail
+
+class TestVscMail(TestCase):
+
+    @mock.patch('vsc.utils.mail.smtplib')
+    @mock.patch('vsc.utils.mail.ssl')
+    def test_send(self, mock_ssl, mock_smtplib):
+
+        msg = MIMEText("test")
+        msg['Subject'] = "subject"
+        msg['From'] = "test@noreply.com"
+        msg['To'] = "test@noreply.com"
+        msg['Reply-to'] = "test@noreply.com"
+
+        vm = VscMail()
+
+        self.assertEqual(vm.mail_host, '')
+        self.assertEqual(vm.mail_port, 0)
+        self.assertEqual(vm.smtp_auth_user, None)
+        self.assertEqual(vm.smtp_auth_password, None)
+        self.assertEqual(vm.smtp_use_starttls, False)
+
+        vm._send(mail_from="test@noreply.com", mail_to="test@noreply.com", mail_subject="s", msg=msg)
+
+        vm = VscMail(
+            mail_host = "test.machine.com",
+            mail_port=123,
+            smtp_auth_user="me",
+            smtp_auth_password="hunter2",
+        )
+
+        self.assertEqual(vm.mail_host, "test.machine.com")
+        self.assertEqual(vm.mail_port, 123)
+        self.assertEqual(vm.smtp_auth_user, "me")
+        self.assertEqual(vm.smtp_auth_password, "hunter2")
+        self.assertEqual(vm.smtp_use_starttls, False)
+
+        vm._send(mail_from="test@noreply.com", mail_to="test@noreply.com", mail_subject="s", msg=msg)
+
+        mock_smtplib.SMTP.assert_called_with(host="test.machine.com", port=123)
+
+        vm = VscMail(
+            mail_host = "test.machine.com",
+            mail_port=124,
+            smtp_auth_user="me",
+            smtp_auth_password="hunter2",
+            smtp_use_starttls=True
+        )
+
+        self.assertEqual(vm.mail_host, "test.machine.com")
+        self.assertEqual(vm.mail_port, 124)
+        self.assertEqual(vm.smtp_auth_user, "me")
+        self.assertEqual(vm.smtp_auth_password, "hunter2")
+        self.assertEqual(vm.smtp_use_starttls, True)
+
+        vm._send(mail_from="test@noreply.com", mail_to="test@noreply.com", mail_subject="s", msg=msg)
+
+        mock_smtplib.SMTP.assert_called_with(host="test.machine.com", port=124)
+        mock_ssl.create_default_context.assert_called()
