Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3import io
4import socket
5import ssl
6import typing
8from ..exceptions import ProxySchemeUnsupported
10if typing.TYPE_CHECKING:
11 from typing_extensions import Self
13 from .ssl_ import _TYPE_PEER_CERT_RET, _TYPE_PEER_CERT_RET_DICT
16_WriteBuffer = typing.Union[bytearray, memoryview]
17_ReturnValue = typing.TypeVar("_ReturnValue")
19SSL_BLOCKSIZE = 16384
22class SSLTransport:
23 """
24 The SSLTransport wraps an existing socket and establishes an SSL connection.
26 Contrary to Python's implementation of SSLSocket, it allows you to chain
27 multiple TLS connections together. It's particularly useful if you need to
28 implement TLS within TLS.
30 The class supports most of the socket API operations.
31 """
33 @staticmethod
34 def _validate_ssl_context_for_tls_in_tls(ssl_context: ssl.SSLContext) -> None:
35 """
36 Raises a ProxySchemeUnsupported if the provided ssl_context can't be used
37 for TLS in TLS.
39 The only requirement is that the ssl_context provides the 'wrap_bio'
40 methods.
41 """
43 if not hasattr(ssl_context, "wrap_bio"):
44 raise ProxySchemeUnsupported(
45 "TLS in TLS requires SSLContext.wrap_bio() which isn't "
46 "available on non-native SSLContext"
47 )
49 def __init__(
50 self,
51 socket: socket.socket,
52 ssl_context: ssl.SSLContext,
53 server_hostname: str | None = None,
54 suppress_ragged_eofs: bool = True,
55 ) -> None:
56 """
57 Create an SSLTransport around socket using the provided ssl_context.
58 """
59 self.incoming = ssl.MemoryBIO()
60 self.outgoing = ssl.MemoryBIO()
62 self.suppress_ragged_eofs = suppress_ragged_eofs
63 self.socket = socket
65 self.sslobj = ssl_context.wrap_bio(
66 self.incoming, self.outgoing, server_hostname=server_hostname
67 )
69 # Perform initial handshake.
70 self._ssl_io_loop(self.sslobj.do_handshake)
72 def __enter__(self) -> Self:
73 return self
75 def __exit__(self, *_: typing.Any) -> None:
76 self.close()
78 def fileno(self) -> int:
79 return self.socket.fileno()
81 def read(self, len: int = 1024, buffer: typing.Any | None = None) -> int | bytes:
82 return self._wrap_ssl_read(len, buffer)
84 def recv(self, buflen: int = 1024, flags: int = 0) -> int | bytes:
85 if flags != 0:
86 raise ValueError("non-zero flags not allowed in calls to recv")
87 return self._wrap_ssl_read(buflen)
89 def recv_into(
90 self,
91 buffer: _WriteBuffer,
92 nbytes: int | None = None,
93 flags: int = 0,
94 ) -> None | int | bytes:
95 if flags != 0:
96 raise ValueError("non-zero flags not allowed in calls to recv_into")
97 if nbytes is None:
98 nbytes = len(buffer)
99 return self.read(nbytes, buffer)
101 def sendall(self, data: bytes, flags: int = 0) -> None:
102 if flags != 0:
103 raise ValueError("non-zero flags not allowed in calls to sendall")
104 count = 0
105 with memoryview(data) as view, view.cast("B") as byte_view:
106 amount = len(byte_view)
107 while count < amount:
108 v = self.send(byte_view[count:])
109 count += v
111 def send(self, data: bytes, flags: int = 0) -> int:
112 if flags != 0:
113 raise ValueError("non-zero flags not allowed in calls to send")
114 return self._ssl_io_loop(self.sslobj.write, data)
116 def makefile(
117 self,
118 mode: str,
119 buffering: int | None = None,
120 *,
121 encoding: str | None = None,
122 errors: str | None = None,
123 newline: str | None = None,
124 ) -> typing.BinaryIO | typing.TextIO | socket.SocketIO:
125 """
126 Python's httpclient uses makefile and buffered io when reading HTTP
127 messages and we need to support it.
129 This is unfortunately a copy and paste of socket.py makefile with small
130 changes to point to the socket directly.
131 """
132 if not set(mode) <= {"r", "w", "b"}:
133 raise ValueError(f"invalid mode {mode!r} (only r, w, b allowed)")
135 writing = "w" in mode
136 reading = "r" in mode or not writing
137 assert reading or writing
138 binary = "b" in mode
139 rawmode = ""
140 if reading:
141 rawmode += "r"
142 if writing:
143 rawmode += "w"
144 raw = socket.SocketIO(self, rawmode) # type: ignore[arg-type]
145 self.socket._io_refs += 1 # type: ignore[attr-defined]
146 if buffering is None:
147 buffering = -1
148 if buffering < 0:
149 buffering = io.DEFAULT_BUFFER_SIZE
150 if buffering == 0:
151 if not binary:
152 raise ValueError("unbuffered streams must be binary")
153 return raw
154 buffer: typing.BinaryIO
155 if reading and writing:
156 buffer = io.BufferedRWPair(raw, raw, buffering) # type: ignore[assignment]
157 elif reading:
158 buffer = io.BufferedReader(raw, buffering)
159 else:
160 assert writing
161 buffer = io.BufferedWriter(raw, buffering)
162 if binary:
163 return buffer
164 text = io.TextIOWrapper(buffer, encoding, errors, newline)
165 text.mode = mode # type: ignore[misc]
166 return text
168 def unwrap(self) -> None:
169 self._ssl_io_loop(self.sslobj.unwrap)
171 def close(self) -> None:
172 self.socket.close()
174 @typing.overload
175 def getpeercert(
176 self, binary_form: typing.Literal[False] = ...
177 ) -> _TYPE_PEER_CERT_RET_DICT | None: ...
179 @typing.overload
180 def getpeercert(self, binary_form: typing.Literal[True]) -> bytes | None: ...
182 def getpeercert(self, binary_form: bool = False) -> _TYPE_PEER_CERT_RET:
183 return self.sslobj.getpeercert(binary_form) # type: ignore[return-value]
185 def version(self) -> str | None:
186 return self.sslobj.version()
188 def cipher(self) -> tuple[str, str, int] | None:
189 return self.sslobj.cipher()
191 def selected_alpn_protocol(self) -> str | None:
192 return self.sslobj.selected_alpn_protocol()
194 def shared_ciphers(self) -> list[tuple[str, str, int]] | None:
195 return self.sslobj.shared_ciphers()
197 def compression(self) -> str | None:
198 return self.sslobj.compression()
200 def settimeout(self, value: float | None) -> None:
201 self.socket.settimeout(value)
203 def gettimeout(self) -> float | None:
204 return self.socket.gettimeout()
206 def _decref_socketios(self) -> None:
207 self.socket._decref_socketios() # type: ignore[attr-defined]
209 def _wrap_ssl_read(self, len: int, buffer: bytearray | None = None) -> int | bytes:
210 try:
211 return self._ssl_io_loop(self.sslobj.read, len, buffer)
212 except ssl.SSLError as e:
213 if e.errno == ssl.SSL_ERROR_EOF and self.suppress_ragged_eofs:
214 if buffer is not None:
215 # eof, no bytes were written into the buffer.
216 return 0
217 # eof, no more data to return.
218 return b""
219 else:
220 raise
222 # func is sslobj.do_handshake or sslobj.unwrap
223 @typing.overload
224 def _ssl_io_loop(self, func: typing.Callable[[], None]) -> None: ...
226 # func is sslobj.write, arg1 is data
227 @typing.overload
228 def _ssl_io_loop(self, func: typing.Callable[[bytes], int], arg1: bytes) -> int: ...
230 # func is sslobj.read, arg1 is len, arg2 is buffer
231 @typing.overload
232 def _ssl_io_loop(
233 self,
234 func: typing.Callable[[int, bytearray | None], bytes],
235 arg1: int,
236 arg2: bytearray | None,
237 ) -> bytes: ...
239 def _ssl_io_loop(
240 self,
241 func: typing.Callable[..., _ReturnValue],
242 arg1: None | bytes | int = None,
243 arg2: bytearray | None = None,
244 ) -> _ReturnValue:
245 """Performs an I/O loop between incoming/outgoing and the socket."""
246 should_loop = True
247 ret = None
249 while should_loop:
250 errno = None
251 try:
252 if arg1 is None and arg2 is None:
253 ret = func()
254 elif arg2 is None:
255 ret = func(arg1)
256 else:
257 ret = func(arg1, arg2)
258 except ssl.SSLError as e:
259 if e.errno not in (ssl.SSL_ERROR_WANT_READ, ssl.SSL_ERROR_WANT_WRITE):
260 # WANT_READ, and WANT_WRITE are expected, others are not.
261 raise e
262 errno = e.errno
264 buf = self.outgoing.read()
265 self.socket.sendall(buf)
267 if errno is None:
268 should_loop = False
269 elif errno == ssl.SSL_ERROR_WANT_READ:
270 buf = self.socket.recv(SSL_BLOCKSIZE)
271 if buf:
272 self.incoming.write(buf)
273 else:
274 self.incoming.write_eof()
275 return typing.cast(_ReturnValue, ret)