Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion massive/rest/base.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import certifi
import json
import os
import urllib3
import inspect
from urllib3.util.retry import Retry
Expand All @@ -21,6 +22,14 @@
pass


def _default_ca_bundle() -> str:
for env_var in ("SSL_CERT_FILE", "REQUESTS_CA_BUNDLE"):
bundle = os.environ.get(env_var)
if bundle:
return bundle
return certifi.where()


class BaseClient:
def __init__(
self,
Expand Down Expand Up @@ -76,7 +85,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=_default_ca_bundle(),
cert_reqs="CERT_REQUIRED",
retries=retry_strategy, # use the customized Retry instance
timeout=self.timeout, # set timeout for each request
Expand Down
4 changes: 2 additions & 2 deletions massive/websocket/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,13 @@
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 ..logging import get_logger
import logging
from ..exceptions import AuthError
from ..rest.base import _default_ca_bundle

env_key = "MASSIVE_API_KEY"
logger = get_logger("WebSocketClient")
Expand Down Expand Up @@ -98,7 +98,7 @@ async def connect(
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.load_verify_locations(_default_ca_bundle())

last_exc = None
async for s in connect(
Expand Down
49 changes: 49 additions & 0 deletions test_rest/test_ssl_ca_bundle.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
import os
import unittest
from unittest import mock

import certifi

from massive import RESTClient
from massive.rest.base import _default_ca_bundle


class SSLCaBundleTest(unittest.TestCase):
def tearDown(self):
for var in ("SSL_CERT_FILE", "REQUESTS_CA_BUNDLE"):
os.environ.pop(var, None)

def test_ssl_cert_file_takes_precedence(self):
with mock.patch.dict(
os.environ,
{
"SSL_CERT_FILE": "/tmp/custom-ca.pem",
"REQUESTS_CA_BUNDLE": "/tmp/other-ca.pem",
},
):
self.assertEqual(_default_ca_bundle(), "/tmp/custom-ca.pem")
client = RESTClient("key")
self.assertEqual(
client.client.connection_pool_kw["ca_certs"], "/tmp/custom-ca.pem"
)

def test_requests_ca_bundle_fallback(self):
os.environ.pop("SSL_CERT_FILE", None)
with mock.patch.dict(os.environ, {"REQUESTS_CA_BUNDLE": "/tmp/other-ca.pem"}):
self.assertEqual(_default_ca_bundle(), "/tmp/other-ca.pem")
client = RESTClient("key")
self.assertEqual(
client.client.connection_pool_kw["ca_certs"], "/tmp/other-ca.pem"
)

def test_certifi_default_without_env(self):
with mock.patch.dict(os.environ, {}, clear=True):
self.assertEqual(_default_ca_bundle(), certifi.where())
client = RESTClient("key")
self.assertEqual(
client.client.connection_pool_kw["ca_certs"], certifi.where()
)


if __name__ == "__main__":
unittest.main()