derivepassphrase.git

@@ -8,7 +8,6 @@ from __future__ import annotations

 import base64

 import collections

-import contextlib

 import copy

 import enum

 import importlib

@@ -16,7 +15,6 @@ import inspect

 import json

 import logging

 import os

-import socket

 import unicodedata

 from typing import (

     TYPE_CHECKING,

@@ -37,6 +35,7 @@ from derivepassphrase import _types, exporter, ssh_agent, vault

 if TYPE_CHECKING:

     import pathlib

+    import socket

     import types

     from collections.abc import (

         Iterator,

@@ -484,21 +483,8 @@ def _get_suitable_ssh_keys(

     Args:

         conn:

-            An optional connection hint to the SSH agent; specifically,

-            an SSH agent client, or a socket connected to an SSH agent.

-

-            If an existing SSH agent client, then this client will be

-            queried for the SSH keys, and otherwise left intact.

-

-            If a socket, then a one-shot client will be constructed

-            based on the socket to query the agent, and deconstructed

-            afterwards.

-

-            If neither are given, then the agent's socket location is

-            looked up in the `SSH_AUTH_SOCK` environment variable, and

-            used to construct/deconstruct a one-shot client, as in the

-            previous case.  This requires the [`socket.AF_UNIX`][]

-            symbol to exist.

+            An optional connection hint to the SSH agent.  See

+            [`ssh_agent.SSHAgentClient.ensure_agent_subcontext`][].

     Yields:

         Every SSH key from the SSH agent that is suitable for passphrase

@@ -524,20 +510,7 @@ def _get_suitable_ssh_keys(

             The agent failed to supply a list of loaded keys.

     """

-    client: ssh_agent.SSHAgentClient

-    client_context: contextlib.AbstractContextManager[Any]

-    # Use match/case here once Python 3.9 becomes unsupported.

-    if isinstance(conn, ssh_agent.SSHAgentClient):

-        client = conn

-        client_context = contextlib.nullcontext()

-    elif isinstance(conn, socket.socket) or conn is None:

-        client = ssh_agent.SSHAgentClient(socket=conn)

-        client_context = client

-    else:  # pragma: no cover

-        assert_never(conn)

-        msg = f'invalid connection hint: {conn!r}'

-        raise TypeError(msg)  # noqa: DOC501

-    with client_context:

+    with ssh_agent.SSHAgentClient.ensure_agent_subcontext(conn) as client:

         try:

             all_key_comment_pairs = list(client.list_keys())

         except EOFError as e:  # pragma: no cover

@@ -633,21 +606,8 @@ def _select_ssh_key(

     Args:

         conn:

-            An optional connection hint to the SSH agent; specifically,

-            an SSH agent client, or a socket connected to an SSH agent.

-

-            If an existing SSH agent client, then this client will be

-            queried for the SSH keys, and otherwise left intact.

-

-            If a socket, then a one-shot client will be constructed

-            based on the socket to query the agent, and deconstructed

-            afterwards.

-

-            If neither are given, then the agent's socket location is

-            looked up in the `SSH_AUTH_SOCK` environment variable, and

-            used to construct/deconstruct a one-shot client, as in the

-            previous case.  This requires the [`socket.AF_UNIX`][]

-            symbol to exist.

+            An optional connection hint to the SSH agent.  See

+            [`ssh_agent.SSHAgentClient.ensure_agent_subcontext`][].

     Returns:

         The selected SSH key.

@@ -7,16 +7,17 @@

 from __future__ import annotations

 import collections

+import contextlib

 import os

 import socket

 from typing import TYPE_CHECKING, overload

-from typing_extensions import Self

+from typing_extensions import Self, assert_never

 from derivepassphrase import _types

 if TYPE_CHECKING:

-    from collections.abc import Iterable, Sequence

+    from collections.abc import Iterable, Iterator, Sequence

     from types import TracebackType

     from typing_extensions import Buffer

@@ -282,6 +283,61 @@ class SSHAgentClient:

             bytes(bytestring[m + HEAD_LEN :]),

         )

+    @classmethod

+    @contextlib.contextmanager

+    def ensure_agent_subcontext(

+        cls,

+        conn: SSHAgentClient | socket.socket | None = None,

+    ) -> Iterator[SSHAgentClient]:

+        """Return an SSH agent client subcontext.

+

+        If necessary, construct an SSH agent client first using the

+        connection hint.

+

+        Args:

+            conn:

+                If an existing SSH agent client, then enter a context

+                within this client's scope.  After exiting the context,

+                the client persists, including its socket.

+

+                If a socket, then construct a client using this socket,

+                then enter a context within this client's scope.  After

+                exiting the context, the client is destroyed and the

+                socket is closed.

+

+                If `None`, construct a client using agent

+                auto-discovery, then enter a context within this

+                client's scope.  After exiting the context, both the

+                client and its socket are destroyed.

+

+        Yields:

+            When entering this context, return the SSH agent client.

+

+        Raises:

+            KeyError:

+                `conn` was `None`, and the `SSH_AUTH_SOCK` environment

+                variable was not found.

+            NotImplementedError:

+                `conn` was `None`, and this Python does not support

+                [`socket.AF_UNIX`][], so the SSH agent client cannot be

+                automatically set up.

+            OSError:

+                `conn` was a socket or `None`, and there was an error

+                setting up a socket connection to the agent.

+

+        """

+        # Use match/case here once Python 3.9 becomes unsupported.

+        if isinstance(conn, SSHAgentClient):

+            with contextlib.nullcontext():

+                yield conn

+        elif isinstance(conn, socket.socket) or conn is None:

+            with SSHAgentClient(socket=conn) as client:

+                yield client

+        else:  # pragma: no cover

+            assert_never(conn)

+            msg = f'invalid connection hint: {conn!r}'

+            raise TypeError(msg)  # noqa: DOC501

+

     @overload

     def request(  # pragma: no cover

         self,

@@ -522,10 +522,10 @@ class Vault:

                 'signature not deterministic'

             )

             raise ValueError(msg)

-        with ssh_agent.SSHAgentClient() as client:

+        with ssh_agent.SSHAgentClient.ensure_agent_subcontext() as client:

             raw_sig = client.sign(key, cls._UUID)

-        _keytype, trailer = client.unstring_prefix(raw_sig)

-        signature_blob = client.unstring(trailer)

+        _keytype, trailer = ssh_agent.SSHAgentClient.unstring_prefix(raw_sig)

+        signature_blob = ssh_agent.SSHAgentClient.unstring(trailer)

         return bytes(base64.standard_b64encode(signature_blob))

     @staticmethod