Merge pull request #211 from KeepSafe/Issue_206_ssl_context

Add ssl_context to TCPConnector

Author: asvetlov

aiohttp/connector.py

@@ -223,10 +223,17 @@ class TCPConnector(BaseConnector):
     """
 
     def __init__(self, *args, verify_ssl=True,
-                 resolve=False, family=socket.AF_INET, **kwargs):
+                 resolve=False, family=socket.AF_INET, ssl_context=None,
+                 **kwargs):
+        if not verify_ssl and ssl_context is not None:
+            raise ValueError(
+                "Either disable ssl certificate validation by "
+                "verify_ssl=False or specify ssl_context, not both.")
+
         super().__init__(*args, **kwargs)
 
         self._verify_ssl = verify_ssl
+        self._ssl_context = ssl_context
         self._family = family
         self._resolve = resolve
         self._resolved_hosts = {}
@@ -236,6 +243,33 @@ def verify_ssl(self):
         """Do check for ssl certifications?"""
         return self._verify_ssl
 
+    @property
+    def ssl_context(self):
+        """SSLContext instance for https requests.
+
+        Lazy property, creates context on demand.
+        """
+        if self._ssl_context is None:
+            if not self._verify_ssl:
+                sslcontext = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
+                sslcontext.options |= ssl.OP_NO_SSLv2
+                sslcontext.options |= ssl.OP_NO_SSLv3
+                sslcontext.options |= getattr(ssl, "OP_NO_COMPRESSION", 0)
+                sslcontext.set_default_verify_paths()
+            elif hasattr(ssl, 'create_default_context'):
+                # Python 3.4+
+                sslcontext = ssl.create_default_context()
+            else:  # pragma: no cover
+                # Fallback for Python 3.3.
+                sslcontext = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
+                sslcontext.options |= ssl.OP_NO_SSLv2
+                sslcontext.options |= ssl.OP_NO_SSLv3
+                sslcontext.options |= getattr(ssl, "OP_NO_COMPRESSION", 0)
+                sslcontext.set_default_verify_paths()
+                sslcontext.verify_mode = ssl.CERT_REQUIRED
+            self._ssl_context = sslcontext
+        return self._ssl_context
+
     @property
     def family(self):
         """Socket family like AF_INET."""
@@ -288,11 +322,10 @@ def _create_connection(self, req, **kwargs):
 
         Has same keyword arguments as BaseEventLoop.create_connection.
         """
-        sslcontext = req.ssl
-        if req.ssl and not self._verify_ssl:
-            sslcontext = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
-            sslcontext.options |= ssl.OP_NO_SSLv2
-            sslcontext.set_default_verify_paths()
+        if req.ssl:
+            sslcontext = self.ssl_context
+        else:
+            sslcontext = None
 
         hosts = yield from self._resolve_host(req.host, req.port)
 

tests/test_connector.py

@@ -6,6 +6,7 @@
 import time
 import socket
 import unittest
+import ssl
 from unittest import mock
 
 import aiohttp
@@ -312,6 +313,22 @@ def test_tcp_connector_clear_resolved_hosts(self):
         conn.clear_resolved_hosts()
         self.assertEqual(conn.resolved_hosts, {})
 
+    def test_ambigous_verify_ssl_and_ssl_context(self):
+        with self.assertRaises(ValueError):
+            aiohttp.TCPConnector(
+                verify_ssl=False,
+                ssl_context=ssl.SSLContext(ssl.PROTOCOL_SSLv23))
+
+    def test_dont_recreate_ssl_context(self):
+        conn = aiohttp.TCPConnector(loop=self.loop)
+        ctx = conn.ssl_context
+        self.assertIs(ctx, conn.ssl_context)
+
+    def test_respect_precreated_ssl_context(self):
+        ctx = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
+        conn = aiohttp.TCPConnector(loop=self.loop, ssl_context=ctx)
+        self.assertIs(ctx, conn.ssl_context)
+
 
 class HttpClientConnectorTests(unittest.TestCase):