-
Notifications
You must be signed in to change notification settings - Fork 34
Description
The Python ssl module documentation says:
Since Python 3.2 and 2.7.9, it is recommended to use the SSLContext.wrap_socket() of an SSLContext instance to wrap sockets as SSLSocket objects. The helper functions create_default_context() returns a new context with secure default settings. The old wrap_socket() function is deprecated since it is both inefficient and has no support for server name indication (SNI) and hostname matching.
It'd be nice if sslpsk followed suit and provided a similar API. It seems that 3.7+ provides some core API improvements that would make this possible. I played around and this seemed to work:
def _ssl_setup_psk_callbacks(sslobj):
psk = sslobj.context.psk
hint = sslobj.context.hint
if psk:
if sslobj.server_side:
cb = psk if callable(psk) else lambda _identity: psk
_ssl_set_psk_server_callback(sslobj, cb, hint)
else:
cb = psk if callable(psk) else lambda _hint: psk if isinstance(psk, tuple) else (psk, b"")
_ssl_set_psk_client_callback(sslobj, cb)
class SSLPSKContext(ssl.SSLContext):
@property
def psk(self):
return getattr(self, "_psk", None)
@psk.setter
def psk(self, psk):
self._psk = psk
@property
def hint(self):
return getattr(self, "_hint", None)
@hint.setter
def hint(self, hint):
self._hint = hint
class SSLPSKObject(ssl.SSLObject):
def do_handshake(self, *args, **kwargs):
_ssl_setup_psk_callbacks(self)
super().do_handshake(*args, **kwargs)
class SSLPSKSocket(ssl.SSLSocket):
def do_handshake(self, *args, **kwargs):
_ssl_setup_psk_callbacks(self)
super().do_handshake(*args, **kwargs)
SSLPSKContext.sslobject_class = SSLPSKObject
SSLPSKContext.sslsocket_class = SSLPSKSocket(It seems like SSLPSKSocket alone is useful to replace the existing functionality. I think SSLPSKObject is useful only under certain frameworks. I could not test that, so perhaps it'd be smarter to even avoid including that for now.)
With that, one can use SSLPSKContext where they'd use SSLContext before, and SSLPSKContext.psk = … to set the PSK (and .hint = … for the hint).
A backwards-compatible sslpsk.wrap_socket() can still be offered with:
def wrap_socket(sock, psk, hint=None,
server_side=False,
ssl_version=ssl.PROTOCOL_TLS,
do_handshake_on_connect=True,
suppress_ragged_eofs=True,
ciphers=None):
context = SSLPSKContext(ssl_version)
if ciphers:
context.set_ciphers(ciphers)
context.psk = psk
context.hint = hint
return context.wrap_socket(
sock=sock, server_side=server_side,
do_handshake_on_connect=do_handshake_on_connect,
suppress_ragged_eofs=suppress_ragged_eofs
)