Add an SSLContextAdapter

AUTHORS.rst

@@ -50,3 +50,5 @@ Patches and Suggestions
 
 - Achim Herwig <[email protected]>
 
+- Nikos Graser <[email protected]>
+

docs/adapters.rst

@@ -21,6 +21,8 @@ with requests. The transport adapters are all kept in
 
 - :class:`requests_toolbelt.adapters.host_header_ssl.HostHeaderSSLAdapter`
 
+- :class:`requests_toolbelt.adapters.ssl_context.SSLContextAdapter`
+
 AppEngineAdapter
 ----------------
 
@@ -243,3 +245,19 @@ specifically for that domain, instead of adding it to every ``https://`` and
 
 .. autoclass:: requests_toolbelt.adapters.socket_options.TCPKeepAliveAdapter
 
+SSLContextAdapter
+-----------------
+
+.. note::
+
+    This adapter will only work with requests 2.4.0 or newer. The ability to
+    pass arbitrary ssl contexts does not exist prior to requests 2.4.0.
+
+The ``SSLContextAdapter`` allows the user to pass an arbitrary SSLContext
+object from Python's ``ssl`` library that will be used for all connections
+made through it.
+
+While not suitable for general-purpose usage, this allows more control over
+the SSL-related behaviour of ``requests``.
+
+.. autoclass:: requests_toolbelt.adapters.ssl_context.SSLContextAdapter

requests_toolbelt/adapters/ssl_context.py

@@ -0,0 +1,53 @@
+# -*- coding: utf-8 -*-
+"""This file contains an implementation of the SSLContextAdapter.
+
+   It requires a version of requests >= 2.4.0.
+"""
+
+import requests
+from requests.adapters import HTTPAdapter
+
+
+class SSLContextAdapter(HTTPAdapter):
+    """An adapter that lets the user inject a custom SSL context for all
+       requests made through it.
+
+       The SSL context will simply be passed through to urllib3, which
+       causes it to skip creation of its own context.
+
+       Note that the SSLContext is not persisted when pickling - this is on
+       purpose.
+       So, after unpickling the SSLContextAdapter will behave like an
+       HTTPAdapter until a new SSLContext is set.
+
+       Example usage:
+
+       .. code-block:: python
+
+           import requests
+           from ssl import create_default_context
+           from requests import Session
+           from requests_toolbelt.adapters.ssl_context import SSLContextAdapter
+
+           s = Session()
+           s.mount('https://', SSLContextAdapter(ssl_context=create_default_context()))
+    """
+
+    def __init__(self, **kwargs):
+        self.ssl_context = None
+        if 'ssl_context' in kwargs:
+            self.ssl_context = kwargs['ssl_context']
+            del kwargs['ssl_context']
+
+        super(SSLContextAdapter, self).__init__(**kwargs)
+
+    def __setstate__(self, state):
+        # SSLContext objects aren't picklable and shouldn't be persisted anyway
+        self.ssl_context = None
+        super(SSLContextAdapter, self).__setstate__(state)
+
+    def init_poolmanager(self, *args, **kwargs):
+        if requests.__build__ >= 0x020400:
+            if 'ssl_context' not in kwargs:
+                kwargs['ssl_context'] = self.ssl_context
+        super(SSLContextAdapter, self).init_poolmanager(*args, **kwargs)

tests/test_ssl_context_adapter.py

@@ -0,0 +1,60 @@
+# -*- coding: utf-8 -*-
+"""Tests for the SSLContextAdapter."""
+
+import pickle
+from ssl import SSLContext, PROTOCOL_TLSv1
+
+import mock
+import pytest
+import requests
+from requests.adapters import HTTPAdapter
+from requests_toolbelt.adapters.ssl_context import SSLContextAdapter
+
+
+@pytest.mark.skipif(requests.__build__ < 0x020400,
+                    reason="Test case for newer requests versions.")
+@mock.patch.object(HTTPAdapter, 'init_poolmanager')
+def test_ssl_context_arg_is_passed_on_newer_requests(init_poolmanager):
+    ssl_context = SSLContext(PROTOCOL_TLSv1)
+    SSLContextAdapter(
+        pool_connections=10,
+        pool_maxsize=5,
+        max_retries=0,
+        pool_block=True,
+        ssl_context=ssl_context
+    )
+    init_poolmanager.assert_called_once_with(
+        10, 5, block=True, ssl_context=ssl_context
+    )
+
+
+@pytest.mark.skipif(requests.__build__ >= 0x020400,
+                    reason="Test case for older requests versions.")
+@mock.patch.object(HTTPAdapter, 'init_poolmanager')
+def test_ssl_context_arg_is_not_passed_on_older_requests(init_poolmanager):
+    ssl_context = SSLContext(PROTOCOL_TLSv1)
+    SSLContextAdapter(
+        pool_connections=10,
+        pool_maxsize=5,
+        max_retries=0,
+        pool_block=True,
+        ssl_context=ssl_context
+    )
+    init_poolmanager.assert_called_once_with(
+        10, 5, block=True
+    )
+
+
+def test_adapter_has_ssl_context_attr():
+    ssl_context = SSLContext(PROTOCOL_TLSv1)
+    adapter = SSLContextAdapter(ssl_context=ssl_context)
+
+    assert adapter.ssl_context is ssl_context
+
+
+def test_adapter_loses_ssl_context_after_pickling():
+    ssl_context = SSLContext(PROTOCOL_TLSv1)
+    adapter = SSLContextAdapter(ssl_context=ssl_context)
+    adapter = pickle.loads(pickle.dumps(adapter))
+
+    assert adapter.ssl_context is None