Repository: abhinavsingh/proxy.py
Pull request #1443 (2024-08-09 17:28:29)

Deprecate usage of `ssl.wrap_socket` in favour of `SSLContext.wrap_socket`


Patch:
diff --git a/proxy/common/constants.py b/proxy/common/constants.py
index 673f9a9..4655822 100644
--- a/proxy/common/constants.py
+++ b/proxy/common/constants.py
@@ -9,6 +9,7 @@
     :license: BSD, see LICENSE for more details.
 """
 import os
+import ssl
 import sys
 import time
 import pathlib
@@ -156,6 +157,9 @@ DEFAULT_ENABLE_PROXY_PROTOCOL = False
 DEFAULT_SELECTOR_SELECT_TIMEOUT = 25 / 1000
 DEFAULT_WAIT_FOR_TASKS_TIMEOUT = 1 / 1000
 DEFAULT_INACTIVE_CONN_CLEANUP_TIMEOUT = 1   # in seconds
+DEFAULT_SSL_CONTEXT_OPTIONS = (
+    ssl.OP_NO_SSLv2 | ssl.OP_NO_SSLv3 | ssl.OP_NO_TLSv1 | ssl.OP_NO_TLSv1_1
+)
 
 DEFAULT_DEVTOOLS_DOC_URL = 'http://proxy'
 DEFAULT_DEVTOOLS_FRAME_ID = secrets.token_hex(8)
diff --git a/proxy/common/utils.py b/proxy/common/utils.py
index 4e2eed8..9ac8883 100644
--- a/proxy/common/utils.py
+++ b/proxy/common/utils.py
@@ -26,7 +26,7 @@ from typing import Any, Dict, List, Type, Tuple, Callable, Optional
 from .types import HostPort
 from .constants import (
     CRLF, COLON, HTTP_1_1, IS_WINDOWS, WHITESPACE, DEFAULT_TIMEOUT,
-    DEFAULT_THREADLESS, PROXY_AGENT_HEADER_VALUE,
+    DEFAULT_THREADLESS, PROXY_AGENT_HEADER_VALUE, DEFAULT_SSL_CONTEXT_OPTIONS,
 )
 
 
@@ -219,20 +219,11 @@ def wrap_socket(
         cafile: Optional[str] = None,
 ) -> ssl.SSLSocket:
     """Use this to upgrade server_side socket to TLS."""
-    ctx = ssl.create_default_context(
-        ssl.Purpose.CLIENT_AUTH,
-        cafile=cafile,
-    )
-    ctx.options |= ssl.OP_NO_SSLv2 | ssl.OP_NO_SSLv3 | ssl.OP_NO_TLSv1 | ssl.OP_NO_TLSv1_1
+    ctx = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH, cafile=cafile)
+    ctx.options |= DEFAULT_SSL_CONTEXT_OPTIONS
     ctx.verify_mode = ssl.CERT_NONE
-    ctx.load_cert_chain(
-        certfile=certfile,
-        keyfile=keyfile,
-    )
-    return ctx.wrap_socket(
-        conn,
-        server_side=True,
-    )
+    ctx.load_cert_chain(certfile=certfile, keyfile=keyfile)
+    return ctx.wrap_socket(conn, server_side=True)
 
 
 def new_socket_connection(
diff --git a/proxy/core/connection/client.py b/proxy/core/connection/client.py
index f241c56..4c8ad97 100644
--- a/proxy/core/connection/client.py
+++ b/proxy/core/connection/client.py
@@ -9,11 +9,12 @@
     :license: BSD, see LICENSE for more details.
 """
 import ssl
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from .types import tcpConnectionTypes
 from .connection import TcpConnection, TcpConnectionUninitializedException
 from ...common.types import HostPort, TcpOrTlsSocket
+from ...common.constants import DEFAULT_SSL_CONTEXT_OPTIONS
 
 
 class TcpClientConnection(TcpConnection):
@@ -42,11 +43,19 @@ class TcpClientConnection(TcpConnection):
     def wrap(self, keyfile: str, certfile: str) -> None:
         self.connection.setblocking(True)
         self.flush()
-        self._conn = ssl.wrap_socket(
-            self.connection,
-            server_side=True,
-            certfile=certfile,
-            keyfile=keyfile,
-            ssl_version=ssl.PROTOCOL_TLS,
+        ctx = ssl.SSLContext(
+            protocol=(
+                ssl.PROTOCOL_TLS_CLIENT
+                if self.tag == 'server'
+                else ssl.PROTOCOL_TLS_SERVER
+            ),
         )
+        ctx.options |= DEFAULT_SSL_CONTEXT_OPTIONS
+        ctx.load_cert_chain(certfile=certfile, keyfile=keyfile)
+        assert self.addr
+        kwargs: Dict[str, Any] = {'server_side': True}
+        if self.tag == 'server':
+            assert self.addr
+            kwargs['server_hostname'] = self.addr[0]
+        self._conn = ctx.wrap_socket(self.connection, **kwargs)
         self.connection.setblocking(False)


Test patch:
diff --git a/tests/http/proxy/test_http_proxy_tls_interception.py b/tests/http/proxy/test_http_proxy_tls_interception.py
index 654bbc5..2fbdaef 100644
--- a/tests/http/proxy/test_http_proxy_tls_interception.py
+++ b/tests/http/proxy/test_http_proxy_tls_interception.py
@@ -62,9 +62,9 @@ class TestHttpProxyTlsInterception(Assertions):
         self.mock_ssl_context.return_value.wrap_socket.return_value = upstream_tls_sock
 
         # Used for client wrapping
-        self.mock_ssl_wrap = mocker.patch('ssl.wrap_socket')
+        self.mock_ssl_wrap = mocker.patch('ssl.SSLContext')
         client_tls_sock = mock.MagicMock(spec=ssl.SSLSocket)
-        self.mock_ssl_wrap.return_value = client_tls_sock
+        self.mock_ssl_wrap.return_value.wrap_socket.return_value = client_tls_sock
 
         plain_connection = mock.MagicMock(spec=socket.socket)
 
@@ -251,13 +251,18 @@ class TestHttpProxyTlsInterception(Assertions):
         )
         assert self.flags.ca_cert_dir is not None
         self.mock_ssl_wrap.assert_called_with(
-            self._conn,
-            server_side=True,
+            protocol=ssl.PROTOCOL_TLS_SERVER,
+        )
+        self.mock_ssl_wrap.return_value.load_cert_chain(
             keyfile=self.flags.ca_signing_key_file,
             certfile=HttpProxyPlugin.generated_cert_file_path(
-                self.flags.ca_cert_dir, host,
+                self.flags.ca_cert_dir,
+                host,
             ),
-            ssl_version=ssl.PROTOCOL_TLS,
+        )
+        self.mock_ssl_wrap.return_value.wrap_socket.assert_called_with(
+            self._conn,
+            server_side=True,
         )
         self.assertEqual(self._conn.setblocking.call_count, 2)
         self.assertEqual(
diff --git a/tests/plugin/test_http_proxy_plugins_with_tls_interception.py b/tests/plugin/test_http_proxy_plugins_with_tls_interception.py
index 3d8d6a2..a0a05b6 100644
--- a/tests/plugin/test_http_proxy_plugins_with_tls_interception.py
+++ b/tests/plugin/test_http_proxy_plugins_with_tls_interception.py
@@ -46,7 +46,7 @@ class TestHttpProxyPluginExamplesWithTlsInterception(Assertions):
             'proxy.http.proxy.server.TcpServerConnection',
         )
         self.mock_ssl_context = mocker.patch('ssl.create_default_context')
-        self.mock_ssl_wrap = mocker.patch('ssl.wrap_socket')
+        self.mock_ssl_wrap = mocker.patch('ssl.SSLContext')
 
         self.mock_sign_csr.return_value = True
         self.mock_gen_csr.return_value = True
@@ -82,7 +82,7 @@ class TestHttpProxyPluginExamplesWithTlsInterception(Assertions):
         self.server_ssl_connection = mocker.MagicMock(spec=ssl.SSLSocket)
         self.mock_ssl_context.return_value.wrap_socket.return_value = self.server_ssl_connection
         self.client_ssl_connection = mocker.MagicMock(spec=ssl.SSLSocket)
-        self.mock_ssl_wrap.return_value = self.client_ssl_connection
+        self.mock_ssl_wrap.return_value.wrap_socket.return_value = self.client_ssl_connection
 
         def has_buffer() -> bool:
             return cast(bool, self.server.queue.called)