diff --git a/massive/certs.py b/massive/certs.py new file mode 100644 index 00000000..5a6b4cb9 --- /dev/null +++ b/massive/certs.py @@ -0,0 +1,17 @@ +import os + +import certifi + + +def ca_bundle_path() -> str: + """Path of the CA bundle used to verify TLS connections. + + Honors the standard ``SSL_CERT_FILE`` and ``REQUESTS_CA_BUNDLE`` + environment variables (used for example behind corporate TLS + inspection), falling back to the certifi bundle. + """ + return ( + os.environ.get("SSL_CERT_FILE") + or os.environ.get("REQUESTS_CA_BUNDLE") + or certifi.where() + ) diff --git a/massive/rest/base.py b/massive/rest/base.py index 3349d7ef..958a0b44 100644 --- a/massive/rest/base.py +++ b/massive/rest/base.py @@ -1,4 +1,3 @@ -import certifi import json import urllib3 import inspect @@ -8,6 +7,7 @@ from datetime import datetime from importlib.metadata import version, PackageNotFoundError from .models.request import RequestOptionBuilder +from ..certs import ca_bundle_path from ..logging import get_logger import logging from urllib.parse import urlencode, urlparse @@ -76,7 +76,7 @@ def __init__( self.client = urllib3.PoolManager( num_pools=num_pools, headers=self.headers, # default headers sent with each request. - ca_certs=certifi.where(), + ca_certs=ca_bundle_path(), cert_reqs="CERT_REQUIRED", retries=retry_strategy, # use the customized Retry instance timeout=self.timeout, # set timeout for each request diff --git a/massive/websocket/__init__.py b/massive/websocket/__init__.py index 0d5409cd..bd936a2c 100644 --- a/massive/websocket/__init__.py +++ b/massive/websocket/__init__.py @@ -5,10 +5,10 @@ import json import asyncio import ssl -import certifi from .models import * from websockets.asyncio.client import connect, ClientConnection from websockets.exceptions import ConnectionClosedOK, ConnectionClosedError +from ..certs import ca_bundle_path from ..logging import get_logger import logging from ..exceptions import AuthError @@ -75,6 +75,14 @@ def __init__( else: self.json = json + def _ssl_context(self) -> Optional[ssl.SSLContext]: + # darwin needs some extra <3 + if not self.url.startswith("wss://"): + return None + context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + context.load_verify_locations(ca_bundle_path()) + return context + # https://websockets.readthedocs.io/en/stable/reference/client.html#opening-a-connection async def connect( self, @@ -94,11 +102,7 @@ async def connect( """ reconnects = 0 logger.debug("connect: %s", self.url) - # darwin needs some extra <3 - ssl_context = None - if self.url.startswith("wss://"): - ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) - ssl_context.load_verify_locations(certifi.where()) + ssl_context = self._ssl_context() last_exc = None async for s in connect( diff --git a/test_rest/test_ssl.py b/test_rest/test_ssl.py new file mode 100644 index 00000000..503065c3 --- /dev/null +++ b/test_rest/test_ssl.py @@ -0,0 +1,58 @@ +import os +import unittest +from unittest import mock + +import certifi + +from massive import RESTClient +from massive.certs import ca_bundle_path + + +class CaBundlePathTest(unittest.TestCase): + def test_ssl_cert_file_env_wins(self): + with mock.patch.dict( + os.environ, {"SSL_CERT_FILE": "/custom/ca.pem"}, clear=True + ): + self.assertEqual(ca_bundle_path(), "/custom/ca.pem") + + def test_requests_ca_bundle_env_used(self): + with mock.patch.dict( + os.environ, {"REQUESTS_CA_BUNDLE": "/custom/requests-ca.pem"}, clear=True + ): + self.assertEqual(ca_bundle_path(), "/custom/requests-ca.pem") + + def test_ssl_cert_file_preferred_over_requests_ca_bundle(self): + with mock.patch.dict( + os.environ, + { + "SSL_CERT_FILE": "/custom/ca.pem", + "REQUESTS_CA_BUNDLE": "/custom/requests-ca.pem", + }, + clear=True, + ): + self.assertEqual(ca_bundle_path(), "/custom/ca.pem") + + def test_defaults_to_certifi(self): + with mock.patch.dict(os.environ, {}, clear=True): + self.assertEqual(ca_bundle_path(), certifi.where()) + + +class RestClientSslTest(unittest.TestCase): + def _pool_ca_certs(self): + return RESTClient("test-key").client.connection_pool_kw["ca_certs"] + + def test_rest_client_uses_ssl_cert_file(self): + with mock.patch.dict( + os.environ, {"SSL_CERT_FILE": "/custom/ca.pem"}, clear=True + ): + self.assertEqual(self._pool_ca_certs(), "/custom/ca.pem") + + def test_rest_client_uses_requests_ca_bundle(self): + with mock.patch.dict( + os.environ, {"REQUESTS_CA_BUNDLE": "/custom/requests-ca.pem"}, clear=True + ): + self.assertEqual(self._pool_ca_certs(), "/custom/requests-ca.pem") + + def test_rest_client_defaults_to_certifi(self): + with mock.patch.dict(os.environ, {}, clear=True): + self.assertEqual(self._pool_ca_certs(), certifi.where()) diff --git a/test_websocket/test_ssl.py b/test_websocket/test_ssl.py new file mode 100644 index 00000000..a3ce6f06 --- /dev/null +++ b/test_websocket/test_ssl.py @@ -0,0 +1,45 @@ +import os +import ssl +import unittest +from unittest import mock + +import certifi + +from massive import WebSocketClient + + +class WebSocketSslTest(unittest.TestCase): + def _client(self, secure=True): + return WebSocketClient(api_key="test-key", secure=secure) + + def _loaded_bundle(self, client, env): + with mock.patch.dict(os.environ, env, clear=True): + with mock.patch.object( + ssl.SSLContext, "load_verify_locations", autospec=True + ) as load: + context = client._ssl_context() + self.assertIsNotNone(context) + return load.call_args[0][0] + + def test_wss_uses_ssl_cert_file(self): + self.assertEqual( + self._loaded_bundle(self._client(), {"SSL_CERT_FILE": "/custom/ca.pem"}), + "/custom/ca.pem", + ) + + def test_wss_uses_requests_ca_bundle(self): + self.assertEqual( + self._loaded_bundle( + self._client(), {"REQUESTS_CA_BUNDLE": "/custom/requests-ca.pem"} + ), + "/custom/requests-ca.pem", + ) + + def test_wss_defaults_to_certifi(self): + self.assertEqual( + self._loaded_bundle(self._client(), {}), + certifi.where(), + ) + + def test_plain_ws_has_no_ssl_context(self): + self.assertIsNone(self._client(secure=False)._ssl_context())