1 files changed, 42 insertions, 28 deletions

diff --git a/Lib/smtplib.py b/Lib/smtplib.py
index 9080e4a..b89b5a5 100644
--- a/Lib/smtplib.py
+++ b/Lib/smtplib.py
@@ -173,27 +173,6 @@ try:
 except ImportError:
     _have_ssl = False
 else:
-    class SSLFakeFile:
-        """A fake file like object that really wraps a SSLObject.
-
-        It only supports what is needed in smtplib.
-        """
-        def __init__(self, sslobj):
-            self.sslobj = sslobj
-
-        def readline(self):
-            str = b""
-            chr = None
-            while chr != b"\n":
-                chr = self.sslobj.read(1)
-                if not chr:
-                    break
-                str += chr
-            return str
-
-        def close(self):
-            pass
-
     _have_ssl = True
 
 
@@ -270,6 +249,19 @@ class SMTP:
                     pass
                 self.local_hostname = '[%s]' % addr
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        try:
+            code, message = self.docmd("QUIT")
+            if code != 221:
+                raise SMTPResponseException(code, message)
+        except SMTPServerDisconnected:
+            pass
+        finally:
+            self.close()
+
     def set_debuglevel(self, debuglevel):
         """Set the debug output level.
 
@@ -310,6 +302,7 @@ class SMTP:
         if self.debuglevel > 0:
             print('connect:', (host, port), file=stderr)
         self.sock = self._get_socket(host, port, self.timeout)
+        self.file = None
         (code, msg) = self.getreply()
         if self.debuglevel > 0:
             print("connect:", msg, file=stderr)
@@ -623,7 +616,7 @@ class SMTP:
         # We could not login sucessfully. Return result of last attempt.
         raise SMTPAuthenticationError(code, resp)
 
-    def starttls(self, keyfile=None, certfile=None):
+    def starttls(self, keyfile=None, certfile=None, context=None):
         """Puts the connection to the SMTP server into TLS mode.
 
         If there has been no previous EHLO or HELO command this session, this
@@ -647,8 +640,17 @@ class SMTP:
         if resp == 220:
             if not _have_ssl:
                 raise RuntimeError("No SSL support included in this Python")
-            self.sock = ssl.wrap_socket(self.sock, keyfile, certfile)
-            self.file = SSLFakeFile(self.sock)
+            if context is not None and keyfile is not None:
+                raise ValueError("context and keyfile arguments are mutually "
+                                 "exclusive")
+            if context is not None and certfile is not None:
+                raise ValueError("context and certfile arguments are mutually "
+                                 "exclusive")
+            if context is not None:
+                self.sock = context.wrap_socket(self.sock)
+            else:
+                self.sock = ssl.wrap_socket(self.sock, keyfile, certfile)
+            self.file = None
             # RFC 3207:
             # The client MUST discard any knowledge obtained from
             # the server, such as the list of SMTP service extensions,
@@ -828,24 +830,35 @@ if _have_ssl:
         support). If host is not specified, '' (the local host) is used. If port is
         omitted, the standard SMTP-over-SSL port (465) is used. keyfile and certfile
         are also optional - they can contain a PEM formatted private key and
-        certificate chain file for the SSL connection.
+        certificate chain file for the SSL connection. context also optional, can contain
+        a SSLContext, and is an alternative to keyfile and certfile; If it is specified both
+        keyfile and certfile must be None.
         """
 
         default_port = SMTP_SSL_PORT
 
         def __init__(self, host='', port=0, local_hostname=None,
                      keyfile=None, certfile=None,
-                     timeout=socket._GLOBAL_DEFAULT_TIMEOUT):
+                     timeout=socket._GLOBAL_DEFAULT_TIMEOUT, context=None):
+            if context is not None and keyfile is not None:
+                raise ValueError("context and keyfile arguments are mutually "
+                                 "exclusive")
+            if context is not None and certfile is not None:
+                raise ValueError("context and certfile arguments are mutually "
+                                 "exclusive")
             self.keyfile = keyfile
             self.certfile = certfile
+            self.context = context
             SMTP.__init__(self, host, port, local_hostname, timeout)
 
         def _get_socket(self, host, port, timeout):
             if self.debuglevel > 0:
                 print('connect:', (host, port), file=stderr)
             new_socket = socket.create_connection((host, port), timeout)
-            new_socket = ssl.wrap_socket(new_socket, self.keyfile, self.certfile)
-            self.file = SSLFakeFile(new_socket)
+            if self.context is not None:
+                new_socket = self.context.wrap_socket(new_socket)
+            else:
+                new_socket = ssl.wrap_socket(new_socket, self.keyfile, self.certfile)
             return new_socket
 
     __all__.append("SMTP_SSL")
@@ -882,6 +895,7 @@ class LMTP(SMTP):
         # Handle Unix-domain sockets.
         try:
             self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+            self.file = None
             self.sock.connect(host)
         except socket.error as msg:
             if self.debuglevel > 0: