gh-63284: Add support for TLS-PSK (pre-shared key) to the ssl module

Author: grantramsay

Doc/library/ssl.rst

@@ -1987,6 +1987,79 @@ to speed up repeated connections from the same clients.
          >>> ssl.create_default_context().verify_mode  # doctest: +SKIP
          <VerifyMode.CERT_REQUIRED: 2>
 
+.. method:: SSLContext.set_psk_client_callback(callback)
+
+   Enables TLS-PSK (pre-shared key) authentication on a client-side connection.
+
+   In general, certificate based authentication should be preferred over this method.
+
+   The parameter ``callback`` is a callable object with the signature:
+   ``def callback(hint: str | None) -> tuple[str | None, bytes]``.
+   The ``hint`` parameter is an optional identity hint sent by the server.
+   The return value is a tuple in the form (client-identity, psk).
+   Client-identity is an optional string which may be used by the server to
+   select a corresponding PSK for the client. PSK is a
+   :term:`bytes-like object` represening the pre-shared key.
+
+   Setting ``callback`` to :const:`None` removes any existing callback.
+
+   Example usage::
+
+      context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+      context.check_hostname = False
+      context.verify_mode = ssl.CERT_NONE
+      context.maximum_version = ssl.TLSVersion.TLSv1_2
+      context.set_ciphers('PSK')
+
+      # A simple lambda:
+      psk = bytes.fromhex('deadbeef')
+      context.set_psk_client_callback(lambda hint: (None, psk))
+
+      # A table using the hint from the server:
+      psk_table = { 'ServerId_1': bytes.fromhex('deadbeef'),
+                    'ServerId_2': bytes.fromhex('cafebabe')
+      }
+      def callback(hint):
+          return 'ClientId_1', psk_table[hint]
+      context.set_psk_client_callback(callback)
+
+   .. versionadded:: 3.12
+
+.. method:: SSLContext.set_psk_server_callback(callback, identity_hint=None)
+
+   Enables TLS-PSK (pre-shared key) authentication on a server-side connection.
+
+   In general, certificate based authentication should be preferred over this method.
+
+   The parameter ``callback`` is a callable object with the signature:
+   ``def callback(identity: str | None) -> bytes``.
+   The ``identity`` parameter is an optional identity sent by the client which can
+   be used to select a corresponding PSK.
+   The return value is a :term:`bytes-like object` represening the pre-shared key.
+
+   Setting ``callback`` to :const:`None` removes any existing callback.
+
+   The parameter ``identity_hint`` is an optional identity hint sent to the client.
+
+   Example usage::
+
+      context.maximum_version = ssl.TLSVersion.TLSv1_2
+      context.set_ciphers('PSK')
+
+      # A simple lambda:
+      psk = bytes.fromhex('deadbeef')
+      context.set_psk_server_callback(lambda identity: psk)
+
+      # A table using the identity of the client:
+      psk_table = { 'ClientId_1': bytes.fromhex('deadbeef'),
+                    'ClientId_2': bytes.fromhex('cafed00d')
+      }
+      def callback(identity):
+          return psk_table[identity]
+      context.set_psk_server_callback(callback, 'ServerId_1')
+
+   .. versionadded:: 3.12
+
 .. index:: single: certificates
 
 .. index:: single: X509 certificate

Include/internal/pycore_global_objects_fini_generated.h

@@ -301,6 +301,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(call)
         STRUCT_FOR_ID(call_exception_handler)
         STRUCT_FOR_ID(call_soon)
+        STRUCT_FOR_ID(callback)
         STRUCT_FOR_ID(cancel)
         STRUCT_FOR_ID(capath)
         STRUCT_FOR_ID(category)
@@ -440,6 +441,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(hook)
         STRUCT_FOR_ID(id)
         STRUCT_FOR_ID(ident)
+        STRUCT_FOR_ID(identity_hint)
         STRUCT_FOR_ID(ignore)
         STRUCT_FOR_ID(imag)
         STRUCT_FOR_ID(importlib)

Include/internal/pycore_global_strings.h

@@ -4203,6 +4203,61 @@ def test_session_handling(self):
                 self.assertEqual(str(e.exception),
                                  'Session refers to a different SSLContext.')
 
+    def test_psk(self):
+        psk = bytes.fromhex('deadbeef')
+
+        client_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
+        client_context.check_hostname = False
+        client_context.verify_mode = ssl.CERT_NONE
+        client_context.maximum_version = ssl.TLSVersion.TLSv1_2
+        client_context.set_ciphers('PSK')
+        client_context.set_psk_client_callback(lambda hint: (None, psk))
+
+        server_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+        server_context.maximum_version = ssl.TLSVersion.TLSv1_2
+        server_context.set_ciphers('PSK')
+        server_context.set_psk_server_callback(lambda identity: psk)
+
+        # correct PSK should connect
+        server = ThreadedEchoServer(context=server_context)
+        with server:
+            with client_context.wrap_socket(socket.socket()) as s:
+                s.connect((HOST, server.port))
+
+        # incorrect PSK should fail
+        incorrect_psk = bytes.fromhex('cafebabe')
+        client_context.set_psk_client_callback(lambda hint: (None, incorrect_psk))
+        server = ThreadedEchoServer(context=server_context)
+        with server:
+            with client_context.wrap_socket(socket.socket()) as s:
+                with self.assertRaises(ssl.SSLError):
+                    s.connect((HOST, server.port))
+
+        # identity_hint and client_identity should be sent to the other side
+        identity_hint = 'identity-hint'
+        client_identity = 'client-identity'
+
+        def client_callback(hint):
+            self.assertEqual(hint, identity_hint)
+            return client_identity, psk
+
+        def server_callback(identity):
+            self.assertEqual(identity, client_identity)
+            return psk
+
+        client_context.set_psk_client_callback(client_callback)
+        server_context.set_psk_server_callback(server_callback, identity_hint)
+        server = ThreadedEchoServer(context=server_context)
+        with server:
+            with client_context.wrap_socket(socket.socket()) as s:
+                s.connect((HOST, server.port))
+
+        # adding client callback to server or vice versa raises an exception
+        with self.assertRaisesRegex(ssl.SSLError, 'Cannot add PSK server callback'):
+            client_context.set_psk_server_callback(server_callback, identity_hint)
+        with self.assertRaisesRegex(ssl.SSLError, 'Cannot add PSK client callback'):
+            server_context.set_psk_client_callback(client_callback)
+
 
 @unittest.skipUnless(has_tls_version('TLSv1_3'), "Test needs TLS 1.3")
 class TestPostHandshakeAuth(unittest.TestCase):

Include/internal/pycore_runtime_init_generated.h

@@ -1466,6 +1466,7 @@ Ajith Ramachandran
 Dhushyanth Ramasamy
 Ashwin Ramaswami
 Jeff Ramnani
+Grant Ramsay
 Bayard Randel
 Varpu Rantala
 Brodie Rao

Include/internal/pycore_unicodeobject_generated.h

@@ -0,0 +1 @@
+Added support for TLS-PSK (pre-shared key) to the ssl module