derivepassphrase.git

@@ -687,7 +687,7 @@ class SSH_AGENTC(int, enum.Enum):  # noqa: N801

     EXTENSION = 27

     """"""

-    def __bytes__(self) -> bytes:

+    def __bytes__(self) -> bytes:  # pragma: no cover [unused]

         """Return the corresponding request code."""

         return int.to_bytes(self, 1, "big", signed=False)

@@ -724,7 +724,7 @@ class SSH_AGENT(int, enum.Enum):  # noqa: N801

     EXTENSION_RESPONSE = 29

     """"""

-    def __bytes__(self) -> bytes:

+    def __bytes__(self) -> bytes:  # pragma: no cover [unused]

         """Return the corresponding response code."""

         return int.to_bytes(self, 1, "big", signed=False)

@@ -308,12 +308,15 @@ class AgentProtocolResponseQueue:

     ([`AgentProtocolRequest`][data.AgentProtocolRequest]) to multiple

     response messages

     ([`AgentProtocolResponse`][data.AgentProtocolResponse]).  Whenever

-    [`_request_message`][] is called with a payload that matches

+    [`_send_request_message`][] is called with a payload that matches

     a [`AgentProtocolRequest`][data.AgentProtocolRequest] from the map,

     the corresponding

     [`AgentProtocolResponse`][data.AgentProtocolResponse] objects are

-    added to the queue, to be read by subsequent [`_response_pair`][]

-    calls.

+    added to the queue, to be read by subsequent

+    [`_get_response_pair`][] calls.  The [`_send_request_message`][] and

+    [`_get_response_pair`][] methods are function- and

+    interface-compatible with the internal methods of the same name in

+    [`ssh_agent.SSHAgentClient`][].

     The queue is a context manager.  When entered or exited, it

     verifies that the queue is empty.

@@ -390,13 +393,37 @@ class AgentProtocolResponseQueue:

     _types.BuiltinSSHAgentSocketProvider.STUB_AGENT

 )

 class StubbedSSHAgentSocket:

-    """A stubbed SSH agent presenting an [`_types.SSHAgentSocket`][]."""

+    """A stubbed SSH agent presenting an [`_types.SSHAgentSocket`][].

+

+    On the network protocol side, the agent implements the full

+    [`_types.SSHAgentSocket`][] interface, including pipelined and

+    unaligned agent requests.  However, on the application side, the

+    agent is intrinsically tied to [the set of SSH test

+    keys][data.ALL_KEYS], and only gives meaningful answers for

+    operations on the test keys and for agent operations in use by

+    [`ssh_agent.SSHAgentClient`][].  The agent does not actually

+    implement any cryptography; all cryptography-related answers are

+    derived from the recorded test key data.

+

+    It is not safe to further monkeypatch the agent's [`recv`][] or

+    [`sendall`][] methods on their own: either monkeypatch both of them,

+    or manipulate the [`send_to_client`][] bytes queue directly instead.

+    Given an [`ssh_agent.SSHAgentClient`][] connected to

+    a [`StubbedSSHAgentSocket`][], if the test ensures proper message

+    serialization and protocol framing and if the monkeypatching can be

+    expressed in terms of full request messages and full response

+    messages, prefer using a [`AgentProtocolResponseQueue`][] to

+    monkeypatch the client's high-level request/response-loop instead of

+    monkeypatching the low-level socket communication in this agent

+    socket.

+

+    """

-    _SOCKET_IS_CLOSED = "Socket is closed."

     _NO_FLAG_SUPPORT = "This stubbed SSH agent socket does not support flags."

     _PROTOCOL_VIOLATION = "SSH agent protocol violation."

     _INVALID_REQUEST = "Invalid request."

     _UNSUPPORTED_REQUEST = "Unsupported request."

+    _INCOMPLETE_REQUEST = "The last request was incomplete."

     HEADER_SIZE = 4

     CODE_SIZE = 1

@@ -435,8 +462,15 @@ class StubbedSSHAgentSocket:

         return self

     def __exit__(self, *args: object) -> None:

-        """Mark the agent's socket as closed."""

+        """Mark the agent's socket as closed.

+

+        Raises:

+            AssertionError:

+                The last request to the agent was incomplete.

+

+        """

         self.closed = True

+        assert not self.receive_from_client, self._INCOMPLETE_REQUEST

     def sendall(self, data: Buffer, flags: int = 0, /) -> None:

         """Send data to the SSH agent.

@@ -458,34 +492,66 @@ class StubbedSSHAgentSocket:

         Raises:

             AssertionError:

                 The flags argument, if specified, must be 0.

-            ValueError:

-                The agent's socket is already closed.  No further

-                requests can be sent.

+            OSError:

+                The socket connection is already closed.

         """

         assert not flags, self._NO_FLAG_SUPPORT

-        if self.closed:

-            raise ValueError(self._SOCKET_IS_CLOSED)

+        self._check_for_io_on_closed_connection()

         self.receive_from_client.extend(memoryview(data))

-        try:

-            self.parse_client_request_and_dispatch()

-        except ValueError:

-            payload = int.to_bytes(_types.SSH_AGENT.FAILURE, 1, "big")

-            self.send_to_client.extend(int.to_bytes(len(payload), 4, "big"))

-            self.send_to_client.extend(payload)

-        finally:

-            self.receive_from_client.clear()

+        while self.receive_from_client:

+            result: Buffer | Iterator[int]

+            if len(self.receive_from_client) < self.HEADER_SIZE:

+                break

+            count = int.from_bytes(

+                self.receive_from_client[: self.HEADER_SIZE],

+                "big",

+                signed=False,

+            )

+            if count:

+                code = int.from_bytes(

+                    self.receive_from_client[

+                        self.HEADER_SIZE : self.HEADER_SIZE + self.CODE_SIZE

+                    ],

+                    "big",

+                    signed=False,

+                )

+                request = bytes(

+                    self.receive_from_client[: self.HEADER_SIZE + count]

+                )

+                if len(request) < self.HEADER_SIZE + count:

+                    break

+                request_payload = request[self.HEADER_SIZE + self.CODE_SIZE :]

+

+                if code == _types.SSH_AGENTC.REQUEST_IDENTITIES:

+                    result = self.request_identities(list_extended=False)

+                elif code == _types.SSH_AGENTC.SIGN_REQUEST:

+                    result = self.sign(request_payload)

+                elif self._check_for_extension(code, "query"):

+                    result = self.query_extensions()

+                elif self._check_for_extension(

+                    code, "list-extended@putty.projects.tartarus.org"

+                ):

+                    result = self.request_identities(list_extended=True)

+                else:

+                    result = self._failure()

+            else:

+                request = bytes(self.receive_from_client[: self.HEADER_SIZE])

+                result = self._failure()

+            self.send_to_client.extend(

+                ssh_agent.SSHAgentClient.string(bytes(result))

+            )

+            self.receive_from_client[: len(request)] = b""

     def recv(self, count: int, flags: int = 0, /) -> bytes:

         """Read data from the SSH agent.

         As per the SSH agent protocol, data is only available to be read

-        immediately after a request via [`sendall`][].  Calls to

-        [`recv`][] at other points in time that attempt to read data

-        violate the protocol, and will fail.  Notwithstanding the last

-        sentence, at any point in time, though pointless, it is

-        additionally permissible to read 0 bytes from the agent, or any

-        number of bytes from a closed socket.

+        immediately after a request via [`sendall`][] and if the socket

+        connection is still open.  Calls to [`recv`][] at other points

+        in time that attempt to read data violate the protocol, and will

+        fail.  (A [`recv`][] of zero bytes does not read data.)  Calls

+        to [`recv`][] when the socket connection is closed always fail.

         Args:

             count:

@@ -495,8 +561,8 @@ class StubbedSSHAgentSocket:

         Returns:

             (A chunk of) the SSH agent's response to the most recent

-            request.  If reading 0 bytes, or if reading from a closed

-            socket, the returned chunk is always an empty byte string.

+            request.  If reading 0 bytes, the returned chunk is always

+            an empty byte string.

         Raises:

             AssertionError:

@@ -505,44 +571,25 @@ class StubbedSSHAgentSocket:

                 Alternatively, `recv` was called when there was no

                 response to be obtained, in violation of the SSH agent

                 protocol.

+            OSError:

+                The socket connection is already closed.

         """

         assert not flags, self._NO_FLAG_SUPPORT

-        assert not count or self.closed or self.send_to_client, (

-            self._PROTOCOL_VIOLATION

-        )

+        self._check_for_io_on_closed_connection()

+        assert not count or self.send_to_client, self._PROTOCOL_VIOLATION

         ret = bytes(self.send_to_client[:count])

         del self.send_to_client[:count]

         return ret

-    def parse_client_request_and_dispatch(self) -> None:

-        """Parse the client request and call the matching handler.

+    def _failure(self) -> bytes:

+        return bytes(_types.SSH_AGENT.FAILURE)

-        This agent supports the

-        [`SSH_AGENTC_REQUEST_IDENTITIES`][_types.SSH_AGENTC.REQUEST_IDENTITIES],

-        [`SSH_AGENTC_SIGN_REQUEST`][_types.SSH_AGENTC.SIGN_REQUEST] and

-        the [`SSH_AGENTC_EXTENSION`][_types.SSH_AGENTC.EXTENSION]

-        request types.

-

-        """

-

-        if len(self.receive_from_client) < self.HEADER_SIZE + self.CODE_SIZE:

-            raise ValueError(self._INVALID_REQUEST)

-        target_header = ssh_agent.SSHAgentClient.uint32(

-            len(self.receive_from_client) - self.HEADER_SIZE

-        )

-        if target_header != self.receive_from_client[: self.HEADER_SIZE]:

-            raise ValueError(self._INVALID_REQUEST)

-        code = _types.SSH_AGENTC(

-            int.from_bytes(

-                self.receive_from_client[

-                    self.HEADER_SIZE : self.HEADER_SIZE + self.CODE_SIZE

-                ],

-                "big",

-            )

-        )

+    def _check_for_io_on_closed_connection(self) -> None:

+        if self.closed:

+            raise OSError(errno.EBADF, os.strerror(errno.EBADF))

-        def is_enabled_extension(extension: str) -> bool:

+    def _check_for_extension(self, code: int, extension: str) -> bool:

         if (

             extension not in self.enabled_extensions

             or code != _types.SSH_AGENTC.EXTENSION

@@ -552,21 +599,6 @@ class StubbedSSHAgentSocket:

         extension_marker = b"\x1b" + string(extension.encode("ascii"))

         return self.receive_from_client.startswith(extension_marker, 4)

-        result: Buffer | Iterator[int]

-        if code == _types.SSH_AGENTC.REQUEST_IDENTITIES:

-            result = self.request_identities(list_extended=False)

-        elif code == _types.SSH_AGENTC.SIGN_REQUEST:

-            result = self.sign()

-        elif is_enabled_extension("query"):

-            result = self.query_extensions()

-        elif is_enabled_extension("list-extended@putty.projects.tartarus.org"):

-            result = self.request_identities(list_extended=True)

-        else:

-            raise ValueError(self._UNSUPPORTED_REQUEST)

-        self.send_to_client.extend(

-            ssh_agent.SSHAgentClient.string(bytes(result))

-        )

-

     def query_extensions(self) -> Iterator[int]:

         """Answer an `SSH_AGENTC_EXTENSION` request.

@@ -633,9 +665,14 @@ class StubbedSSHAgentSocket:

                     ssh_agent.SSHAgentClient.uint32(0)

                 )

-    def sign(self) -> bytes:

+    def sign(self, request_payload: bytes, /) -> bytes:

         """Answer an `SSH_AGENTC_SIGN_REQUEST` request.

+        Args:

+            request_payload:

+                The data of the sign request, without the protocol

+                framing or the request code.

+

         Returns:

             The bytes payload of the response, without the protocol

             framing.

@@ -647,17 +684,20 @@ class StubbedSSHAgentSocket:

         )

         spec = data.SSHTestKeyDeterministicSignatureClass.SPEC

         rfc6979 = data.SSHTestKeyDeterministicSignatureClass.RFC_6979

+        try:

             key_blob, rest = ssh_agent.SSHAgentClient.unstring_prefix(

-            self.receive_from_client[self.HEADER_SIZE + self.CODE_SIZE :]

+                request_payload

             )

             sign_data, rest = ssh_agent.SSHAgentClient.unstring_prefix(rest)

+        except ValueError:

+            return self._failure()

         if len(rest) != 4:

-            raise ValueError(self._INVALID_REQUEST)

+            return self._failure()

         flags = int.from_bytes(rest, "big")

         if flags:

-            raise ValueError(self._UNSUPPORTED_REQUEST)

+            return self._failure()

         if sign_data != vault.Vault.UUID:

-            raise ValueError(self._UNSUPPORTED_REQUEST)

+            return self._failure()

         for key in data.ALL_KEYS.values():

             if key.public_key_data == key_blob:

                 if spec in key.expected_signatures:

@@ -672,8 +712,8 @@ class StubbedSSHAgentSocket:

                     ) + ssh_agent.SSHAgentClient.string(

                         key.expected_signatures[rfc6979].signature

                     )

-                raise ValueError(self._UNSUPPORTED_REQUEST)

-        raise ValueError(self._UNSUPPORTED_REQUEST)

+                return self._failure()

+        return self._failure()

 # Standard variant

@@ -15,12 +15,15 @@ from __future__ import annotations

 import base64

 import contextlib

 import errno

+import math

 import os

 import pathlib

 import re

 from typing import TYPE_CHECKING

+import hypothesis

 import pytest

+from hypothesis import strategies

 from derivepassphrase import _types, ssh_agent, vault

 from tests import data, machinery

@@ -191,7 +194,6 @@ class Parametrize:

         "message",

         [

             pytest.param(b"\x00\x00\x00\x00", id="empty-message"),

-            pytest.param(b"\x00\x00\x00\x0f\x0d", id="truncated-message"),

             pytest.param(

                 b"\x00\x00\x00\x06\x1b\x00\x00\x00\x01\xff",

                 id="invalid-extension-name",

@@ -200,6 +202,10 @@ class Parametrize:

                 b"\x00\x00\x00\x11\x0d\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00",

                 id="sign-with-trailing-data",

             ),

+            pytest.param(

+                b"\x00\x00\x00\x05\x0d\x00\x00\x00\x00",

+                id="sign-without-payload",

+            ),

         ],

     )

     UNSUPPORTED_SSH_AGENT_MESSAGES = pytest.mark.parametrize(

@@ -243,6 +249,8 @@ class Parametrize:

                     ])

                 ),

                 id="sign-key-no-expected-signature",

+                # NOTE: only unsupported when stubbed agent has RFC 6979

+                # support disabled

             ),

             pytest.param(

                 ssh_agent.SSHAgentClient.string(

@@ -264,6 +272,79 @@ class Parametrize:

     )

+class Strategies:

+    """Common hypothesis data generation strategies."""

+

+    @strategies.composite

+    @staticmethod

+    def proper_bytestring_partition(

+        draw: strategies.DrawFn, bs: bytes | bytearray, /

+    ) -> list[bytes]:

+        """Partition a non-empty string into non-empty parts.

+

+        A string of length `n` is partitioned [`math.isqrt(n

+        - 1)`][math.isqrt] many times, such that each substring is

+        non-empty.

+

+        """

+        n = len(bs)

+        hypothesis.assume(n > 0)

+        num_divisions = math.isqrt(n - 1)

+        divisions = draw(

+            strategies.lists(

+                strategies.integers(1, n - 1),

+                min_size=num_divisions,

+                max_size=num_divisions,

+                unique=True,

+            ).map(sorted),

+            label="divisions",

+        )

+        result: list[bytes] = []

+        i = 0

+        for div in divisions:

+            result.append(bytes(bs[i:div]))

+            i = div

+        result.append(bytes(bs[i:n]))

+        return result

+

+    @strategies.composite

+    @staticmethod

+    def truncated_agent_request(

+        draw: strategies.DrawFn, /, min_size: int = 0, max_size: int = 100

+    ) -> bytes:

+        """Generate a private-use SSH agent request with truncated payload.

+

+        The message will adhere to specified size bounds before

+        truncation.  The request message code will be from the private

+        use area.

+

+        Args:

+            min_size:

+                The minimum size of the request payload, after truncation.

+            max_size:

+                The maximum size of the request payload, before truncation.

+

+        """

+        request_code = draw(

+            strategies.integers(240, 255), label="request_code"

+        )

+        full_message_size = draw(

+            strategies.integers(min_size + 1, max_size),

+            label="full_message_size",

+        )

+        truncated_message = draw(

+            strategies.binary(

+                min_size=min_size, max_size=full_message_size - 1

+            )

+        )

+        payload = bytearray(

+            int.to_bytes(1 + full_message_size, 4, "big", signed=False)

+        )

+        payload.extend(int.to_bytes(request_code, 1, "big", signed=False))

+        payload.extend(truncated_message)

+        return bytes(payload)

+

+

 class TestTestKeys:

     """Tests testing the test keys."""

@@ -543,13 +624,10 @@ class TestStubbedSSHAgentSocketProperOperations:

             # payload: SSH string "query"

             b"\x00\x00\x00\x05query"

         )

-        query_response = b""

-        with pytest.raises(

-            ValueError,

-            match=re.escape(machinery.StubbedSSHAgentSocket._SOCKET_IS_CLOSED),

-        ):

+        with pytest.raises(OSError, match=re.escape(os.strerror(errno.EBADF))):

             agent.sendall(query_request)

-        assert agent.recv(100) == query_response

+        with pytest.raises(OSError, match=re.escape(os.strerror(errno.EBADF))):

+            agent.recv(100)

     def test_no_recv_without_sendall(self) -> None:

         """The agent requires a message before sending a response."""

@@ -562,6 +640,43 @@ class TestStubbedSSHAgentSocketProperOperations:

             ):

                 agent.recv(100)

+    @hypothesis.given(

+        query_request_parts=Strategies.proper_bytestring_partition(

+            b"\x00\x00\x00\x0a\x1b\x00\x00\x00\x05query"

+        )

+    )

+    @hypothesis.example(

+        query_request_parts=[

+            b"",

+            b"\x00\x00\x00",

+            b"\x0a",

+            b"\x1b",

+            b"\x00\x00\x00\x05",

+            b"query",

+        ]

+    )

+    def test_piecemeal_sendall(self, query_request_parts: list[bytes]) -> None:

+        """The agent supports receiving messages incrementally."""

+        with machinery.StubbedSSHAgentSocket() as agent:

+            agent.enabled_extensions = frozenset({"query"})

+            for part in query_request_parts:

+                agent.sendall(part)

+            header = agent.recv(agent.HEADER_SIZE)

+            count = int.from_bytes(header, "big", signed=False)

+            payload = agent.recv(count)

+            assert len(payload) == count

+            assert bytes.startswith(

+                payload,

+                (

+                    bytes(_types.SSH_AGENT.SUCCESS),

+                    bytes(_types.SSH_AGENT.EXTENSION_RESPONSE),

+                ),

+            )

+            code = payload[: agent.CODE_SIZE]

+            assert code == bytes(_types.SSH_AGENT.SUCCESS) or code == bytes(

+                _types.SSH_AGENT.EXTENSION_RESPONSE

+            )

+

     @Parametrize.INVALID_SSH_AGENT_MESSAGES

     def test_invalid_ssh_agent_messages(

         self,

@@ -578,6 +693,19 @@ class TestStubbedSSHAgentSocketProperOperations:

             agent.sendall(message)

             assert agent.recv(100) == query_response

+    @hypothesis.given(message=Strategies.truncated_agent_request())

+    @hypothesis.example(message=b"\x00\x00\x00\x0f\xff")

+    def test_truncated_ssh_agent_message(self, message: bytes) -> None:

+        """The agent diagnoses a final truncated request message."""

+        with pytest.raises(  # noqa: SIM117

+            AssertionError,

+            match=re.escape(

+                machinery.StubbedSSHAgentSocket._INCOMPLETE_REQUEST

+            ),

+        ):

+            with machinery.StubbedSSHAgentSocket() as agent:

+                agent.sendall(message)

+

 class TestStubbedSSHAgentSocketSupportedAndUnsupportedFeatures:

     """Test the stubbed SSH agent socket: supported/unsupported features."""