refactor: Split ssl-specific parts of MocketSocket into MocketSSLSocket

Author: betaboon

mocket/inject.py

@@ -81,14 +81,16 @@ def disable() -> None:
         true_inet_pton,
         true_socket,
         true_socketpair,
-        true_ssl_wrap_socket,
         true_urllib3_match_hostname,
-        true_urllib3_ssl_wrap_socket,
-        true_urllib3_wrap_socket,
     )
     from mocket.ssl.context import (
         true_ssl_context,
     )
+    from mocket.ssl.socket import (
+        true_ssl_wrap_socket,
+        true_urllib3_ssl_wrap_socket,
+        true_urllib3_wrap_socket,
+    )
 
     socket.socket = socket.__dict__["socket"] = true_socket
     socket._socketobject = socket.__dict__["_socketobject"] = true_socket

mocket/socket.py

@@ -7,14 +7,11 @@
 import os
 import select
 import socket
-import ssl
-from datetime import datetime, timedelta
 from json.decoder import JSONDecodeError
 from types import TracebackType
 from typing import Any, Type
 
 import urllib3.connection
-import urllib3.util.ssl_
 from typing_extensions import Self
 
 from mocket.compat import decode_from_bytes, encode_to_bytes
@@ -26,7 +23,6 @@
     Address,
     ReadableBuffer,
     WriteableBuffer,
-    _PeerCertRetDictType,
     _RetAddress,
 )
 from mocket.utils import hexdump, hexload
@@ -38,22 +34,7 @@
 true_inet_pton = socket.inet_pton
 true_socket = socket.socket
 true_socketpair = socket.socketpair
-true_ssl_wrap_socket = None
-
 true_urllib3_match_hostname = urllib3.connection.match_hostname
-true_urllib3_ssl_wrap_socket = urllib3.util.ssl_.ssl_wrap_socket
-true_urllib3_wrap_socket = None
-
-with contextlib.suppress(ImportError):
-    # from Py3.12 it's only under SSLContext
-    from ssl import wrap_socket as ssl_wrap_socket
-
-    true_ssl_wrap_socket = ssl_wrap_socket
-
-with contextlib.suppress(ImportError):
-    from urllib3.util.ssl_ import wrap_socket as urllib3_wrap_socket
-
-    true_urllib3_wrap_socket = urllib3_wrap_socket
 
 
 xxh32 = None
@@ -112,9 +93,6 @@ def _hash_request(h, req):
 
 
 class MocketSocket:
-    cipher = lambda s: ("ADH", "AES256", "SHA")
-    compression = lambda s: ssl.OP_NO_COMPRESSION
-
     def __init__(
         self,
         family: socket.AddressFamily | int = socket.AF_INET,
@@ -129,15 +107,10 @@ def __init__(
 
         self._kwargs = kwargs
         self._true_socket = true_socket(family, type, proto)
-        self._truesocket_recording_dir = None
 
         self._buflen = 65536
         self._timeout: float | None = None
 
-        self._secure_socket = False
-        self._did_handshake = False
-        self._sent_non_empty_bytes = False
-
         self._host = None
         self._port = None
         self._address = None
@@ -204,9 +177,6 @@ def settimeout(self, timeout: float | None) -> None:
     def getsockopt(level: int, optname: int, buflen: int | None = None) -> int:
         return socket.SOCK_STREAM
 
-    def do_handshake(self) -> None:
-        self._did_handshake = True
-
     def getpeername(self) -> _RetAddress:
         return self._address
 
@@ -219,32 +189,6 @@ def getblocking(self) -> bool:
     def getsockname(self) -> _RetAddress:
         return true_gethostbyname(self._address[0]), self._address[1]
 
-    def getpeercert(self, binary_form: bool = False) -> _PeerCertRetDictType:
-        if not (self._host and self._port):
-            self._address = self._host, self._port = Mocket._address
-
-        now = datetime.now()
-        shift = now + timedelta(days=30 * 12)
-        return {
-            "notAfter": shift.strftime("%b %d %H:%M:%S GMT"),
-            "subjectAltName": (
-                ("DNS", f"*.{self._host}"),
-                ("DNS", self._host),
-                ("DNS", "*"),
-            ),
-            "subject": (
-                (("organizationName", f"*.{self._host}"),),
-                (("organizationalUnitName", "Domain Control Validated"),),
-                (("commonName", f"*.{self._host}"),),
-            ),
-        }
-
-    def unwrap(self) -> MocketSocket:
-        return self
-
-    def write(self, data: bytes) -> int | None:
-        return self.send(encode_to_bytes(data))
-
     def connect(self, address: Address) -> None:
         self._address = self._host, self._port = address
         Mocket._address = address
@@ -271,33 +215,26 @@ def sendall(self, data, entry=None, *args, **kwargs):
             self.io.truncate()
             self.io.seek(0)
 
-    def read(self, buffersize: int | None = None) -> bytes:
-        rv = self.io.read(buffersize)
-        if rv:
-            self._sent_non_empty_bytes = True
-        if self._did_handshake and not self._sent_non_empty_bytes:
-            raise ssl.SSLWantReadError("The operation did not complete (read)")
-        return rv
-
     def recv_into(
         self,
         buffer: WriteableBuffer,
         buffersize: int | None = None,
         flags: int | None = None,
     ) -> int:
+        data = self.recv(buffersize)
+
         if hasattr(buffer, "write"):
-            return buffer.write(self.read(buffersize))
+            return buffer.write(data)
+
         # buffer is a memoryview
-        data = self.read(buffersize)
-        if data:
-            buffer[: len(data)] = data
+        buffer[: len(data)] = data
         return len(data)
 
     def recv(self, buffersize: int, flags: int | None = None) -> bytes:
         r_fd, _ = Mocket.get_pair((self._host, self._port))
         if r_fd:
             return os.read(r_fd, buffersize)
-        data = self.read(buffersize)
+        data = self.io.read(buffersize)
         if data:
             return data
         # used by Redis mock
@@ -357,12 +294,6 @@ def true_sendall(self, data: ReadableBuffer, *args: Any, **kwargs: Any) -> int:
             host, port = self._host, self._port
             host = true_gethostbyname(host)
 
-            if isinstance(self._true_socket, true_socket) and self._secure_socket:
-                self._true_socket = true_urllib3_ssl_wrap_socket(
-                    self._true_socket,
-                    **self._kwargs,
-                )
-
             with contextlib.suppress(OSError, ValueError):
                 # already connected
                 self._true_socket.connect((host, port))
@@ -400,6 +331,7 @@ def send(
         *args: Any,
         **kwargs: Any,
     ) -> int:  # pragma: no cover
+        data = encode_to_bytes(data)
         entry = self.get_entry(data)
         if not entry or (entry and self._entry != entry):
             kwargs["entry"] = entry

mocket/ssl/context.py

@@ -5,6 +5,7 @@
 from typing import Any
 
 from mocket.socket import MocketSocket
+from mocket.ssl.socket import MocketSSLSocket
 from mocket.types import ReadableBuffer, StrOrBytesPath
 
 true_ssl_context = ssl.SSLContext
@@ -81,18 +82,16 @@ def check_hostname(self, _: bool) -> None:
         self._check_hostname = False
 
     @staticmethod
-    def wrap_socket(sock: MocketSocket, *args: Any, **kwargs: Any) -> MocketSocket:
-        sock._kwargs = kwargs
-        sock._secure_socket = True
-        return sock
+    def wrap_socket(sock: MocketSocket, *args: Any, **kwargs: Any) -> MocketSSLSocket:
+        return MocketSSLSocket._create(sock, *args, **kwargs)
 
     @staticmethod
     def wrap_bio(
         incoming: Any,  # _ssl.MemoryBIO
         outgoing: Any,  # _ssl.MemoryBIO
         server_side: bool = False,
         server_hostname: str | bytes | None = None,
-    ) -> MocketSocket:
-        ssl_obj = MocketSocket()
+    ) -> MocketSSLSocket:
+        ssl_obj = MocketSSLSocket()
         ssl_obj._host = server_hostname
         return ssl_obj

mocket/ssl/socket.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+import contextlib
+import ssl
+from datetime import datetime, timedelta
+from typing import Any
+
+import urllib3.util.ssl_
+
+from mocket.mocket import Mocket
+from mocket.socket import MocketSocket
+from mocket.types import _PeerCertRetDictType
+
+true_ssl_wrap_socket = None
+true_urllib3_ssl_wrap_socket = urllib3.util.ssl_.ssl_wrap_socket
+true_urllib3_wrap_socket = None
+
+with contextlib.suppress(ImportError):
+    # from Py3.12 it's only under SSLContext
+    from ssl import wrap_socket as ssl_wrap_socket
+
+    true_ssl_wrap_socket = ssl_wrap_socket
+
+with contextlib.suppress(ImportError):
+    from urllib3.util.ssl_ import wrap_socket as urllib3_wrap_socket
+
+    true_urllib3_wrap_socket = urllib3_wrap_socket
+
+
+class MocketSSLSocket(MocketSocket):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+
+        self._did_handshake = False
+
+    def read(self, buffersize: int | None = None) -> bytes:
+        rv = self.io.read(buffersize)
+        if rv:
+            self._sent_non_empty_bytes = True
+        if self._did_handshake and not self._sent_non_empty_bytes:
+            raise ssl.SSLWantReadError("The operation did not complete (read)")
+        return rv
+
+    def write(self, data: bytes) -> int | None:
+        return self.send(data)
+
+    def do_handshake(self) -> None:
+        self._did_handshake = True
+
+    def getpeercert(self, binary_form: bool = False) -> _PeerCertRetDictType:
+        if not (self._host and self._port):
+            self._address = self._host, self._port = Mocket._address
+
+        now = datetime.now()
+        shift = now + timedelta(days=30 * 12)
+        return {
+            "notAfter": shift.strftime("%b %d %H:%M:%S GMT"),
+            "subjectAltName": (
+                ("DNS", f"*.{self._host}"),
+                ("DNS", self._host),
+                ("DNS", "*"),
+            ),
+            "subject": (
+                (("organizationName", f"*.{self._host}"),),
+                (("organizationalUnitName", "Domain Control Validated"),),
+                (("commonName", f"*.{self._host}"),),
+            ),
+        }
+
+    def ciper(self) -> tuple[str, str, str]:
+        return ("ADH", "AES256", "SHA")
+
+    def compression(self) -> str | None:
+        return ssl.OP_NO_COMPRESSION
+
+    def unwrap(self) -> MocketSocket:
+        return self
+
+    @classmethod
+    def _create(cls, sock: MocketSocket, *args, **kwargs) -> MocketSSLSocket:
+        ssl_socket = MocketSSLSocket()
+
+        ssl_socket._kwargs = kwargs
+        ssl_socket._true_socket = true_urllib3_ssl_wrap_socket(
+            sock._true_socket,
+            **kwargs,
+        )
+
+        ssl_socket._timeout = sock._timeout
+
+        ssl_socket._host = sock._host
+        ssl_socket._port = sock._port
+        ssl_socket._address = sock._address
+
+        ssl_socket._io = sock._io
+        ssl_socket._entry = sock._entry
+
+        return ssl_socket

tests/test_http.py

@@ -359,12 +359,12 @@ def test_sockets(self):
         sock = socket.socket(address[0], address[1], address[2])
 
         sock.connect(address[-1])
-        sock.write(f"{method} {path} HTTP/1.0\r\n")
-        sock.write(f"Host: {host}\r\n")
-        sock.write("Content-Type: application/json\r\n")
-        sock.write("Content-Length: %d\r\n" % len(data))
-        sock.write("Connection: close\r\n\r\n")
-        sock.write(data)
+        sock.send(f"{method} {path} HTTP/1.0\r\n".encode())
+        sock.send(f"Host: {host}\r\n".encode())
+        sock.send(b"Content-Type: application/json\r\n")
+        sock.send(b"Content-Length: %d\r\n" % len(data))
+        sock.send(b"Connection: close\r\n\r\n")
+        sock.send(data.encode())
         sock.close()
 
         # Proof that worked.