Add support for STARTTLS for Python 2.6+

imapclient/config.py

@@ -24,6 +24,7 @@ def parse_config_file(path):
         username=None,
         password=None,
         ssl='false',
+        starttls='false',
         stream='false',
         oauth='false',
         oauth_token=None,
@@ -48,6 +49,7 @@ def parse_config_file(path):
         host=parser.get(section, 'host'),
         port=port,
         ssl=parser.getboolean(section, 'ssl'),
+        starttls=parser.getboolean(section, 'starttls'),
         stream=parser.getboolean(section, 'stream'),
 
         username=parser.get(section, 'username'),
@@ -89,6 +91,10 @@ def get_oauth2_token(client_id, client_secret, refresh_token):
 def create_client_from_config(conf):
     client = imapclient.IMAPClient(conf.host, port=conf.port,
                                    ssl=conf.ssl, stream=conf.stream)
+
+    if conf.starttls:
+        client.starttls()
+
     if conf.oauth:
         client.oauth_login(conf.oauth_url,
                            conf.oauth_token,

imapclient/imapclient.py

@@ -12,6 +12,12 @@
 from datetime import datetime
 from operator import itemgetter
 
+try:
+    import ssl
+    HAVE_SSL = True
+except ImportError:
+    HAVE_SSL = False
+
 # Confusingly, this module is for OAUTH v1, not v2
 try:
     import oauth2 as oauth_module
@@ -35,11 +41,15 @@
 
 # We also offer the gmail-specific XLIST command...
 if 'XLIST' not in imaplib.Commands:
-  imaplib.Commands['XLIST'] = imaplib.Commands['LIST']
+    imaplib.Commands['XLIST'] = imaplib.Commands['LIST']
 
 # ...and IDLE
 if 'IDLE' not in imaplib.Commands:
-  imaplib.Commands['IDLE'] = imaplib.Commands['APPEND']
+    imaplib.Commands['IDLE'] = imaplib.Commands['APPEND']
+
+# ...and STARTTLS
+if 'STARTTLS' not in imaplib.Commands:
+    imaplib.Commands['STARTTLS'] = ('NONAUTH',)
 
 
 # System flags
@@ -145,6 +155,120 @@ def _create_IMAP4(self, **kwargs):
         ImapClass = self.ssl and imaplib.IMAP4_SSL or imaplib.IMAP4
         return ImapClass(self.host, self.port, **kwargs)
 
+    def _starttls(self, ssl_context=None, **ssl_opts):
+        """Send a STARTTLS command and wrap the socket with SSL."""
+        name = 'STARTTLS'
+
+        if getattr(self._imap, '_tls_established', False):
+            raise self.AbortError('TLS session already established')
+
+        if not self.has_capability(name):
+            raise self.AbortError('TLS not supported by server')
+
+        typ, dat = self._imap._simple_command(name)
+
+        if typ == 'OK':
+            if ssl_context:
+                server_hostname = (self.host
+                    if getattr(ssl, 'HAS_SNI', False) else None)
+                self._imap.sock = ssl_context.wrap_socket(
+                    self._imap.sock, server_hostname=server_hostname)
+                self._imap.file = self._imap.sock.makefile('rb')
+            else:
+                self._imap.sock = ssl.wrap_socket(self._imap.sock, **ssl_opts)
+                self._imap.file = self._imap.sock.makefile('rb')
+
+            self._imap._tls_established = True
+
+            # refresh cached capabilities,
+            # they may have changed with SSL enabled
+            response = self._imap.untagged_responses.pop('CAPABILITY', None)
+            if response:
+                self._save_capabilities(response[0])
+        else:
+            raise self.Error("Couldn't establish TLS session")
+
+        return self._imap._untagged_response(typ, dat, name)
+
+    def starttls(self, ssl_context=None, **ssl_opts):
+        """Switch to an SSL encrypted connection by sending a STARTTLS command.
+
+        The *ssl_context* argument is optional and should be a
+        :py:class:`ssl.SSLContext` object. If no SSL context is given, and no
+        SSL options are given as extra keyword arguments (see below), a default
+        SSL context will be created, if the :py:mod:`ssl` library supports it
+        (Python 2.7.9 and 3.4+). Otherwise a default set of SSL options will be
+        used. In this case no certificate or hostname checking is performed and
+        SSLv2 is used, falling back to SSLv3, if the server doesn't support it.
+
+        SSL options may be passed as keyword arguments, if *ssl_context* is
+        ``None``. The supported keyword arguments are the same as for
+        :py:func:`ssl.wrap_socket` minus *do_handshake_on_connect* and
+        *suppress_ragged_eofs*, which do not make sense in this context. Also,
+        the *ciphers* option is only supported with Python 2.7 and above.
+
+        If you want consistent SSL options accross all supported Python
+        versions, pass them as keyword arguments and don't use *ciphers*, but
+        then you lose support for automatic loading of system default ca
+        certificates and hostname checking.
+
+        With current Python versions it is recommended to create a default SSL
+        context yourself with :py:func:`ssl.create_default_context`, change
+        options as needed and pass them with the *ssl_context* argument.
+
+        Raises :py:exc:`ValueError` if unrecognized or unsupported keyword
+        arguments or both an SSL context and SSL option arguments are passed.
+
+        Raises :py:exc:`ssl.SSLError` when the SSL connection could not be
+        established.
+
+        Raises :py:exc:`Error` if SSL support is not available and
+        :py:exc:`AbortError` if the server does not support STARTTLS or an
+        SSL connection is already established.
+
+        """
+        if not HAVE_SSL:
+            raise self.Error('SSL support missing')
+
+        # Set SSL options, filling in defaults if necessary
+        default_opts = {
+            'keyfile': None,
+            'certfile': None,
+            'server_side': False,
+            'cert_reqs': ssl.CERT_NONE,
+            'ssl_version': ssl.PROTOCOL_SSLv23,
+            'ca_certs': None,
+        }
+
+        if sys.version_info[:2] >= (2, 7):
+            default_opts['ciphers'] = None
+        elif 'ciphers' in ssl_opts:
+            raise ValueError("Setting SSL ciphers requires Python >= 2.7")
+
+        unknown_opts = set(ssl_opts).difference(default_opts)
+
+        if unknown_opts:
+            raise ValueError("Unrecognized SSL option arguments: %s" %
+                ", ".join(unknown_opts))
+
+        if ssl_context and ssl_opts:
+            raise ValueError("No additional keyword arguments allowed if "
+                "'ssl_context' given")
+
+        default_opts.update(ssl_opts)
+
+        # Try to create a default SSL context if none was passed and no
+        # valid SSL option keywords arguments either
+        try:
+            if ssl_context is None and not ssl_opts:
+                ssl_context = ssl._create_stdlib_context()
+        except AttributeError:
+            # when no default context could be created, use ssl_opts
+            pass
+
+        # Now we either have a valid SSL context or (default) ssl_opts
+        return self._starttls(ssl_context, **default_opts)[0]
+
     def login(self, username, password):
         """Login using *username* and *password*, returning the
         server response.

livetest-sample.ini

@@ -1,4 +1,4 @@
-# Sample configuration file for livetest and interact. 
+# Sample configuration file for livetest and interact.
 
 # An INI file like this is used to specify the IMAP account details to
 # run "live" tests. These files are also supported by IMAPClient's
@@ -12,6 +12,7 @@ password = <password>
 # These are optional
 # port = 143
 # ssl = false
+# starttls = false
 
 # oauth = false
 # oauth_url = ...