|
| 1 | +import unittest |
| 2 | +import ssl |
| 3 | +from unittest.mock import patch, MagicMock |
| 4 | +from tableauserverclient import Server |
| 5 | +from tableauserverclient.server.endpoint import Endpoint |
| 6 | +import logging |
| 7 | + |
| 8 | + |
| 9 | +class TestSSLConfig(unittest.TestCase): |
| 10 | + @patch("requests.session") |
| 11 | + @patch("tableauserverclient.server.endpoint.Endpoint.set_parameters") |
| 12 | + def setUp(self, mock_set_parameters, mock_session): |
| 13 | + """Set up test fixtures with mocked session and request validation""" |
| 14 | + # Mock the session |
| 15 | + self.mock_session = MagicMock() |
| 16 | + mock_session.return_value = self.mock_session |
| 17 | + |
| 18 | + # Mock request preparation |
| 19 | + self.mock_request = MagicMock() |
| 20 | + self.mock_session.prepare_request.return_value = self.mock_request |
| 21 | + |
| 22 | + # Create server instance with mocked components |
| 23 | + self.server = Server("http://test") |
| 24 | + |
| 25 | + def test_default_ssl_config(self): |
| 26 | + """Test that by default, no custom SSL context is used""" |
| 27 | + self.assertIsNone(self.server._ssl_context) |
| 28 | + self.assertNotIn("verify", self.server.http_options) |
| 29 | + |
| 30 | + @patch("ssl.create_default_context") |
| 31 | + def test_weak_dh_config(self, mock_create_context): |
| 32 | + """Test that weak DH keys can be allowed when configured""" |
| 33 | + # Setup mock SSL context |
| 34 | + mock_context = MagicMock() |
| 35 | + mock_create_context.return_value = mock_context |
| 36 | + |
| 37 | + # Configure SSL with weak DH |
| 38 | + self.server.configure_ssl(allow_weak_dh=True) |
| 39 | + |
| 40 | + # Verify SSL context was created and configured correctly |
| 41 | + mock_create_context.assert_called_once() |
| 42 | + mock_context.set_dh_parameters.assert_called_once_with(min_key_bits=512) |
| 43 | + |
| 44 | + # Verify context was added to http options |
| 45 | + self.assertEqual(self.server.http_options["verify"], mock_context) |
| 46 | + |
| 47 | + @patch("ssl.create_default_context") |
| 48 | + def test_disable_weak_dh_config(self, mock_create_context): |
| 49 | + """Test that SSL config can be reset to defaults""" |
| 50 | + # Setup mock SSL context |
| 51 | + mock_context = MagicMock() |
| 52 | + mock_create_context.return_value = mock_context |
| 53 | + |
| 54 | + # First enable weak DH |
| 55 | + self.server.configure_ssl(allow_weak_dh=True) |
| 56 | + self.assertIsNotNone(self.server._ssl_context) |
| 57 | + self.assertIn("verify", self.server.http_options) |
| 58 | + |
| 59 | + # Then disable it |
| 60 | + self.server.configure_ssl(allow_weak_dh=False) |
| 61 | + self.assertIsNone(self.server._ssl_context) |
| 62 | + self.assertNotIn("verify", self.server.http_options) |
| 63 | + |
| 64 | + @patch("ssl.create_default_context") |
| 65 | + def test_warning_on_weak_dh(self, mock_create_context): |
| 66 | + """Test that a warning is logged when enabling weak DH keys""" |
| 67 | + logging.getLogger().setLevel(logging.WARNING) |
| 68 | + with self.assertLogs(level="WARNING") as log: |
| 69 | + self.server.configure_ssl(allow_weak_dh=True) |
| 70 | + self.assertTrue( |
| 71 | + any("WARNING: Allowing weak Diffie-Hellman keys" in record for record in log.output), |
| 72 | + "Expected warning about weak DH keys was not logged", |
| 73 | + ) |
| 74 | + |
| 75 | + |
| 76 | +if __name__ == "__main__": |
| 77 | + unittest.main() |
0 commit comments