#15014: Add 'auth' command to implement auth mechanisms and use it in login.

bitdancer · bitdancer · commit 76e13c1c2978 · 2014-07-03T14:47:46.000-04:00
Patch by Milan Oberkirch.
diff --git a/Doc/library/smtplib.rst b/Doc/library/smtplib.rst
@@ -240,8 +240,7 @@ An :class:`SMTP` instance has the following methods:
    the server is stored as the :attr:`ehlo_resp` attribute, :attr:`does_esmtp`
    is set to true or false depending on whether the server supports ESMTP, and
    :attr:`esmtp_features` will be a dictionary containing the names of the
-   SMTP service extensions this server supports, and their
-   parameters (if any).
+   SMTP service extensions this server supports, and their parameters (if any).
 
    Unless you wish to use :meth:`has_extn` before sending mail, it should not be
    necessary to call this method explicitly.  It will be implicitly called by
@@ -291,6 +290,42 @@ An :class:`SMTP` instance has the following methods:
    :exc:`SMTPException`
       No suitable authentication method was found.
 
+   Each of the authentication methods supported by :mod:`smtplib` are tried in
+   turn if they are advertised as supported by the server (see :meth:`auth`
+   for a list of supported authentication methods).
+
+
+.. method:: SMTP.auth(mechanism, authobject)
+
+   Issue an ``SMTP`` ``AUTH`` command for the specified authentication
+   *mechanism*, and handle the challenge response via *authobject*.
+
+   *mechanism* specifies which authentication mechanism is to
+   be used as argument to the ``AUTH`` command; the valid values are
+   those listed in the ``auth`` element of :attr:`esmtp_features`.
+
+   *authobject* must be a callable object taking a single argument:
+
+     data = authobject(challenge)
+
+   It will be called to process the server's challenge response; the
+   *challenge* argument it is passed will be a ``bytes``.  It should return
+   ``bytes`` *data* that will be base64 encoded and sent to the server.
+
+   The ``SMTP`` class provides ``authobjects`` for the ``CRAM-MD5``, ``PLAIN``,
+   and ``LOGIN`` mechanisms; they are named ``SMTP.auth_cram_md5``,
+   ``SMTP.auth_plain``, and ``SMTP.auth_login`` respectively.  They all require
+   that the ``user`` and ``password`` properties of the ``SMTP`` instance are
+   set to appropriate values.
+
+   User code does not normally need to call ``auth`` directly, but can instead
+   call the :meth:`login` method, which will try each of the above mechanisms in
+   turn, in the order listed.  ``auth`` is exposed to facilitate the
+   implementation of authentication methods not (or not yet) supported directly
+   by :mod:`smtplib`.
+
+   .. versionadded:: 3.5
+
 
 .. method:: SMTP.starttls(keyfile=None, certfile=None, context=None)
 
diff --git a/Doc/whatsnew/3.5.rst b/Doc/whatsnew/3.5.rst
@@ -221,6 +221,13 @@ smtpd
   addresses in the :class:`~smtpd.SMTPServer` constructor, and have it
   successfully connect.  (Contributed by Milan Oberkirch in :issue:`14758`.)
 
+smtplib
+-------
+
+* A new :meth:`~smtplib.SMTP.auth` method provides a convenient way to
+  implement custom authentication mechanisms (contributed by Milan Oberkirch in
+  :issue:`15014`).
+
 socket
 ------
 
diff --git a/Lib/smtplib.py b/Lib/smtplib.py
@@ -571,12 +571,60 @@ def ehlo_or_helo_if_needed(self):
                 if not (200 <= code <= 299):
                     raise SMTPHeloError(code, resp)
 
+    def auth(self, mechanism, authobject):
+        """Authentication command - requires response processing.
+
+        'mechanism' specifies which authentication mechanism is to
+        be used - the valid values are those listed in the 'auth'
+        element of 'esmtp_features'.
+
+        'authobject' must be a callable object taking a single argument:
+
+                data = authobject(challenge)
+
+        It will be called to process the server's challenge response; the
+        challenge argument it is passed will be a bytes.  It should return
+        bytes data that will be base64 encoded and sent to the server.
+        """
+
+        mechanism = mechanism.upper()
+        (code, resp) = self.docmd("AUTH", mechanism)
+        # Server replies with 334 (challenge) or 535 (not supported)
+        if code == 334:
+            challenge = base64.decodebytes(resp)
+            response = encode_base64(
+                authobject(challenge).encode('ascii'), eol='')
+            (code, resp) = self.docmd(response)
+            if code in (235, 503):
+                return (code, resp)
+        raise SMTPAuthenticationError(code, resp)
+
+    def auth_cram_md5(self, challenge):
+        """ Authobject to use with CRAM-MD5 authentication. Requires self.user
+        and self.password to be set."""
+        return self.user + " " + hmac.HMAC(
+            self.password.encode('ascii'), challenge, 'md5').hexdigest()
+
+    def auth_plain(self, challenge):
+        """ Authobject to use with PLAIN authentication. Requires self.user and
+        self.password to be set."""
+        return "\0%s\0%s" % (self.user, self.password)
+
+    def auth_login(self, challenge):
+        """ Authobject to use with LOGIN authentication. Requires self.user and
+        self.password to be set."""
+        (code, resp) = self.docmd(
+            encode_base64(self.user.encode('ascii'), eol=''))
+        if code == 334:
+            return self.password
+        raise SMTPAuthenticationError(code, resp)
+
     def login(self, user, password):
         """Log in on an SMTP server that requires authentication.
 
         The arguments are:
-            - user:     The user name to authenticate with.
-            - password: The password for the authentication.
+            - user:         The user name to authenticate with.
+            - password:     The password for the authentication.
 
         If there has been no previous EHLO or HELO command this session, this
         method tries ESMTP EHLO first.
@@ -593,63 +641,40 @@ def login(self, user, password):
                                   found.
         """
 
-        def encode_cram_md5(challenge, user, password):
-            challenge = base64.decodebytes(challenge)
-            response = user + " " + hmac.HMAC(password.encode('ascii'),
-                                              challenge, 'md5').hexdigest()
-            return encode_base64(response.encode('ascii'), eol='')
-
-        def encode_plain(user, password):
-            s = "\0%s\0%s" % (user, password)
-            return encode_base64(s.encode('ascii'), eol='')
-
-        AUTH_PLAIN = "PLAIN"
-        AUTH_CRAM_MD5 = "CRAM-MD5"
-        AUTH_LOGIN = "LOGIN"
-
         self.ehlo_or_helo_if_needed()
-
         if not self.has_extn("auth"):
             raise SMTPException("SMTP AUTH extension not supported by server.")
 
         # Authentication methods the server claims to support
         advertised_authlist = self.esmtp_features["auth"].split()
 
-        # List of authentication methods we support: from preferred to
-        # less preferred methods. Except for the purpose of testing the weaker
-        # ones, we prefer stronger methods like CRAM-MD5:
-        preferred_auths = [AUTH_CRAM_MD5, AUTH_PLAIN, AUTH_LOGIN]
+        # Authentication methods we can handle in our preferred order:
+        preferred_auths = ['CRAM-MD5', 'PLAIN', 'LOGIN']
 
-        # We try the authentication methods the server advertises, but only the
-        # ones *we* support. And in our preferred order.
-        authlist = [auth for auth in preferred_auths if auth in advertised_authlist]
+        # We try the supported authentications in our preferred order, if
+        # the server supports them.
+        authlist = [auth for auth in preferred_auths
+                    if auth in advertised_authlist]
         if not authlist:
             raise SMTPException("No suitable authentication method found.")
 
         # Some servers advertise authentication methods they don't really
         # support, so if authentication fails, we continue until we've tried
         # all methods.
+        self.user, self.password = user, password
         for authmethod in authlist:
-            if authmethod == AUTH_CRAM_MD5:
-                (code, resp) = self.docmd("AUTH", AUTH_CRAM_MD5)
-                if code == 334:
-                    (code, resp) = self.docmd(encode_cram_md5(resp, user, password))
-            elif authmethod == AUTH_PLAIN:
-                (code, resp) = self.docmd("AUTH",
-                    AUTH_PLAIN + " " + encode_plain(user, password))
-            elif authmethod == AUTH_LOGIN:
-                (code, resp) = self.docmd("AUTH",
-                    "%s %s" % (AUTH_LOGIN, encode_base64(user.encode('ascii'), eol='')))
-                if code == 334:
-                    (code, resp) = self.docmd(encode_base64(password.encode('ascii'), eol=''))
-
-            # 235 == 'Authentication successful'
-            # 503 == 'Error: already authenticated'
-            if code in (235, 503):
-                return (code, resp)
-
-        # We could not login sucessfully. Return result of last attempt.
-        raise SMTPAuthenticationError(code, resp)
+            method_name = 'auth_' + authmethod.lower().replace('-', '_')
+            try:
+                (code, resp) = self.auth(authmethod, getattr(self, method_name))
+                # 235 == 'Authentication successful'
+                # 503 == 'Error: already authenticated'
+                if code in (235, 503):
+                    return (code, resp)
+            except SMTPAuthenticationError as e:
+                last_exception = e
+
+        # We could not login successfully.  Return result of last attempt.
+        raise last_exception
 
     def starttls(self, keyfile=None, certfile=None, context=None):
         """Puts the connection to the SMTP server into TLS mode.
diff --git a/Lib/test/test_smtplib.py b/Lib/test/test_smtplib.py
@@ -10,6 +10,7 @@
 import time
 import select
 import errno
+import base64
 
 import unittest
 from test import support, mock_socket
@@ -605,7 +606,8 @@ def testLineTooLong(self):
     'cram-md5': ('TXIUQUBZB21LD2HLCMUUY29TIDG4OWQ0MJ'
                  'KWZGQ4ODNMNDA4NTGXMDRLZWMYZJDMODG1'),
     }
-sim_auth_login_password = 'C29TZXBHC3N3B3JK'
+sim_auth_login_user = 'TXIUQUBZB21LD2HLCMUUY29T'
+sim_auth_plain = 'AE1YLKFAC29TZXDOZXJLLMNVBQBZB21LCGFZC3DVCMQ='
 
 sim_lists = {'list-1':['Mr.A@somewhere.com','Mrs.C@somewhereesle.com'],
              'list-2':['Ms.B@xn--fo-fka.com',],
@@ -659,18 +661,16 @@ def smtp_EXPN(self, arg):
             self.push('550 No access for you!')
 
     def smtp_AUTH(self, arg):
-        if arg.strip().lower()=='cram-md5':
+        mech = arg.strip().lower()
+        if mech=='cram-md5':
             self.push('334 {}'.format(sim_cram_md5_challenge))
-            return
-        mech, auth = arg.split()
-        mech = mech.lower()
-        if mech not in sim_auth_credentials:
+        elif mech not in sim_auth_credentials:
             self.push('504 auth type unimplemented')
             return
-        if mech == 'plain' and auth==sim_auth_credentials['plain']:
-            self.push('235 plain auth ok')
-        elif mech=='login' and auth==sim_auth_credentials['login']:
-            self.push('334 Password:')
+        elif mech=='plain':
+            self.push('334 ')
+        elif mech=='login':
+            self.push('334 ')
         else:
             self.push('550 No access for you!')
 
@@ -818,28 +818,28 @@ def testEXPN(self):
         self.assertEqual(smtp.expn(u), expected_unknown)
         smtp.quit()
 
-    def testAUTH_PLAIN(self):
-        self.serv.add_feature("AUTH PLAIN")
-        smtp = smtplib.SMTP(HOST, self.port, local_hostname='localhost', timeout=15)
-
-        expected_auth_ok = (235, b'plain auth ok')
-        self.assertEqual(smtp.login(sim_auth[0], sim_auth[1]), expected_auth_ok)
-        smtp.close()
-
-    # SimSMTPChannel doesn't fully support LOGIN or CRAM-MD5 auth because they
-    # require a synchronous read to obtain the credentials...so instead smtpd
+    # SimSMTPChannel doesn't fully support AUTH because it requires a
+    # synchronous read to obtain the credentials...so instead smtpd
     # sees the credential sent by smtplib's login method as an unknown command,
     # which results in smtplib raising an auth error.  Fortunately the error
     # message contains the encoded credential, so we can partially check that it
     # was generated correctly (partially, because the 'word' is uppercased in
     # the error message).
 
+    def testAUTH_PLAIN(self):
+        self.serv.add_feature("AUTH PLAIN")
+        smtp = smtplib.SMTP(HOST, self.port, local_hostname='localhost', timeout=15)
+        try: smtp.login(sim_auth[0], sim_auth[1])
+        except smtplib.SMTPAuthenticationError as err:
+            self.assertIn(sim_auth_plain, str(err))
+        smtp.close()
+
     def testAUTH_LOGIN(self):
         self.serv.add_feature("AUTH LOGIN")
         smtp = smtplib.SMTP(HOST, self.port, local_hostname='localhost', timeout=15)
         try: smtp.login(sim_auth[0], sim_auth[1])
         except smtplib.SMTPAuthenticationError as err:
-            self.assertIn(sim_auth_login_password, str(err))
+            self.assertIn(sim_auth_login_user, str(err))
         smtp.close()
 
     def testAUTH_CRAM_MD5(self):
@@ -857,7 +857,23 @@ def testAUTH_multiple(self):
         smtp = smtplib.SMTP(HOST, self.port, local_hostname='localhost', timeout=15)
         try: smtp.login(sim_auth[0], sim_auth[1])
         except smtplib.SMTPAuthenticationError as err:
-            self.assertIn(sim_auth_login_password, str(err))
+            self.assertIn(sim_auth_login_user, str(err))
+        smtp.close()
+
+    def test_auth_function(self):
+        smtp = smtplib.SMTP(HOST, self.port,
+                            local_hostname='localhost', timeout=15)
+        self.serv.add_feature("AUTH CRAM-MD5")
+        smtp.user, smtp.password = sim_auth[0], sim_auth[1]
+        supported = {'CRAM-MD5': smtp.auth_cram_md5,
+                     'PLAIN': smtp.auth_plain,
+                     'LOGIN': smtp.auth_login,
+                    }
+        for mechanism, method in supported.items():
+            try: smtp.auth(mechanism, method)
+            except smtplib.SMTPAuthenticationError as err:
+                self.assertIn(sim_auth_credentials[mechanism.lower()].upper(),
+                              str(err))
         smtp.close()
 
     def test_with_statement(self):
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -103,6 +103,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #15014: Added 'auth' method to smtplib to make implementing auth
+  mechanisms simpler, and used it internally in the login method.
+
 - Issue #21151: Fixed a segfault in the winreg module when ``None`` is passed
   as a ``REG_BINARY`` value to SetValueEx.  Patch by John Ehresman.
 
