src/fakesshagent/machinery.py (cb7012c) - fakesshagent.git

machinery.py

Blame History Raw

# SPDX-FileCopyrightText: 2026 Marco Ricci <software@the13thletter.info>
#
# SPDX-License-Identifier: Zlib

"""Machinery for the fake SSH agent."""

from __future__ import annotations

import errno
import os
import struct
from typing import TYPE_CHECKING, cast

import derivepassphrase_sshagentsocketprovider as d_sasp
from fakesshagent import _types, data

if TYPE_CHECKING:
    from collections.abc import Generator, Iterable

from typing_extensions import Buffer, Self

__all__ = ("StubbedSSHAgentSocket",)

VAULT_UUID = b"e87eb0f4-34cb-46b9-93ad-766c5ab063e7"
uint32_format = struct.Struct(">I")

def string(bstr: Buffer, /) -> bytes:
    """Return the serialized payload as per the SSH agent protocol."""
    payload = memoryview(bstr)
    buffer = bytearray(payload)
    buffer[:0] = bytes(4)
    uint32_format.pack_into(buffer, 0, len(payload))
    return bytes(buffer)

def uint32(i: int, /) -> bytes:
    """Return the serialized integer as per the SSH agent protocol."""
    return uint32_format.pack(i)

def unstring_prefix(buf: Buffer, /) -> tuple[bytes, bytes]:
    """Decode a string encoded as per the SSH agent protocol.

Any data remaining after the encoded string is returned unchanged.

Args:
        buf:
            A binary buffer, beginning with an encoded string.

Returns:
        The decoded string, and the remaining contents of the buffer, as
        a 2-tuple.

Raises:
        ValueError:
            The encoded string is invalid or truncated.

"""
    payload = memoryview(buf)
    pos = 0
    try:
        size = cast("tuple[int]", uint32_format.unpack_from(payload, 0))[0]
    except struct.error as exc:
        msg = f"Invalid or truncated SSH protocol string value: {buf!r}"
        raise ValueError(msg) from exc
    pos += uint32_format.size
    head = bytes(payload[pos : pos + size])
    if len(head) < size:
        msg = f"Invalid or truncated SSH protocol string value: {buf!r}"
        raise ValueError(msg)
    tail = bytes(payload[pos + size :])
    return head, tail

class StubbedSSHAgentSocket:
    """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.

"""

_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

KNOWN_EXTENSIONS = frozenset({
        "query",
        "list-extended@putty.projects.tartarus.org",
    })
    """Known and implemented protocol extensions."""

def __init__(self, *extensions: str) -> None:
        """Initialize the agent."""
        self.send_to_client = bytearray()
        """
        The buffered response to the client, read piecemeal by [`recv`][].
        """
        self.receive_from_client = bytearray()
        """The last request issued by the client."""
        self.closed = False
        """True if the connection is closed, false otherwise."""
        self.enabled_extensions = frozenset(extensions) & self.KNOWN_EXTENSIONS
        """
        Extensions actually enabled in this particular stubbed SSH agent.
        """
        self.try_rfc6979 = False
        """
        Attempt to issue DSA and ECDSA signatures according to RFC 6979?
        """
        self.try_pageant_068_080 = False
        """
        Attempt to issue DSA and ECDSA signatures as per Pageant 0.68–0.80?
        """  # noqa: RUF001

def __enter__(self) -> Self:
        """Return self."""
        return self

def __exit__(self, *args: object) -> None:
        """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.

The signature, and behavior, is identical to
        [`socket.socket.sendall`][].  Upon successful sending, this
        agent will parse the request, call the appropriate handler, and
        buffer the result such that it can be read via [`recv`][], in
        accordance with the SSH agent protocol.

Args:
            data: Binary data to send to the agent.
            flags: Reserved.  Must be 0.

Raises:
            AssertionError:
                The flags argument, if specified, must be 0.
            OSError:
                The socket connection is already closed.

Note:
            The result should be requested via [`recv`][], and
            interpreted in accordance with the SSH agent protocol.

"""
        assert not flags, self._NO_FLAG_SUPPORT
        self._check_for_io_on_closed_connection()
        self.receive_from_client.extend(memoryview(data))
        while self.receive_from_client:
            result: Buffer | Iterable[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(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`][] 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:
                Number of bytes to read from the agent.
            flags:
                Reserved.  Must be 0.

Returns:
            (A chunk of) the SSH agent's response to the most recent
            request.  If reading 0 bytes, the returned chunk is always
            an empty byte string.

Raises:
            AssertionError:
                The flags argument, if specified, must be 0.

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
        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 _failure(self) -> bytes:  # noqa: PLR6301
        return bytes(_types.SSH_AGENT.FAILURE)

def _check_for_io_on_closed_connection(self) -> None:
        if self.closed:
            raise OSError(errno.EBADF, os.strerror(errno.EBADF))

def _check_for_extension(self, code: int, extension: str) -> bool:
        if (
            extension not in self.enabled_extensions
            or code != _types.SSH_AGENTC.EXTENSION
        ):
            return False
        extension_marker = b"\x1b" + string(extension.encode("ascii"))
        return self.receive_from_client.startswith(extension_marker, 4)

def query_extensions(self) -> Generator[int, None, None]:  # noqa: PLR6301
        """Answer an `SSH_AGENTC_EXTENSION` request.

Yields:
            The bytes payload of the response, without the protocol
            framing.  The payload is yielded byte by byte, as an
            iterable of 8-bit integers.

"""
        yield _types.SSH_AGENT.EXTENSION_RESPONSE
        yield from string(b"query")
        extension_answers = [
            b"query",
            b"list-extended@putty.projects.tartarus.org",
        ]
        for a in extension_answers:
            yield from string(a)

def request_identities(
        self, *, list_extended: bool = False
    ) -> Generator[int, None, None]:
        """Answer an `SSH_AGENTC_REQUEST_IDENTITIES` request.

Args:
            list_extended:
                If true, answer an `SSH_AGENTC_EXTENSION` request for
                the `list-extended@putty.projects.tartarus.org`
                extension. Otherwise, answer an
                `SSH_AGENTC_REQUEST_IDENTITIES` request.

Yields:
            The bytes payload of the response, without the protocol
            framing.  The payload is yielded byte by byte, as an
            iterable of 8-bit integers.

"""
        if list_extended:
            yield _types.SSH_AGENT.SUCCESS
        else:
            yield _types.SSH_AGENT.IDENTITIES_ANSWER
        signature_classes = [
            data.SSHTestKeyDeterministicSignatureClass.SPEC,
        ]
        if (
            "list-extended@putty.projects.tartarus.org"
            in self.enabled_extensions
        ):
            signature_classes.append(
                data.SSHTestKeyDeterministicSignatureClass.RFC_6979
            )
        keys = [
            v
            for v in data.ALL_KEYS.values()
            if any(cls in v.expected_signatures for cls in signature_classes)
        ]
        yield from uint32(len(keys))
        for key in keys:
            yield from string(key.public_key_data)
            yield from string(b"test key without passphrase")
            if list_extended:
                yield from string(uint32(0))

def sign(self, request_payload: bytes, /) -> bytes:  # noqa: PLR0911
        """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.

"""
        try_rfc6979 = (
            "list-extended@putty.projects.tartarus.org"
            in self.enabled_extensions
        )
        spec = data.SSHTestKeyDeterministicSignatureClass.SPEC
        rfc6979 = data.SSHTestKeyDeterministicSignatureClass.RFC_6979
        try:
            key_blob, rest = unstring_prefix(request_payload)
            sign_data, rest = unstring_prefix(rest)
        except ValueError:
            return self._failure()
        if len(rest) != uint32_format.size:  # flags are uint32
            return self._failure()
        flags = int.from_bytes(rest, "big")
        if flags:
            return self._failure()
        if sign_data != VAULT_UUID:
            return self._failure()
        for key in data.ALL_KEYS.values():
            if key.public_key_data == key_blob:
                if spec in key.expected_signatures:
                    return int.to_bytes(
                        _types.SSH_AGENT.SIGN_RESPONSE, 1, "big"
                    ) + string(key.expected_signatures[spec].signature)
                if (
                    try_rfc6979 and rfc6979 in key.expected_signatures
                ):  # pragma: no cover [external]
                    return int.to_bytes(
                        _types.SSH_AGENT.SIGN_RESPONSE, 1, "big"
                    ) + string(key.expected_signatures[rfc6979].signature)
                return self._failure()
        return self._failure()

class StubbedSSHAgentSocketWithAddress(StubbedSSHAgentSocket):
    """A [`StubbedSSHAgentSocket`][] requiring a specific address."""

ADDRESS = "stub-ssh-agent:"
    """The correct address for connecting to this stubbed agent."""

def __init__(self, *extensions: str) -> None:
        """Initialize the agent, based on `SSH_AUTH_SOCK`.

Socket addresses of the form `stub-ssh-agent:<errno_value>` will
        raise an [`OSError`][] (or the respective subclass) with the
        specified [`errno`][] value.  For example,
        `stub-ssh-agent:EPERM` will raise a [`PermissionError`][].

Raises:
            KeyError:
                The `SSH_AUTH_SOCK` environment variable is not set.
            OSError:
                The address in `SSH_AUTH_SOCK` is unsuited.

"""
        super().__init__(*extensions)
        try:
            orig_address = os.environ["SSH_AUTH_SOCK"]
        except KeyError as exc:
            msg = "SSH_AUTH_SOCK environment variable"
            raise KeyError(msg) from exc
        address = orig_address
        if not address.startswith(self.ADDRESS):
            address = self.ADDRESS + "ENOENT"
        errcode = address.removeprefix(self.ADDRESS)
        if errcode and not (
            errcode.startswith("E") and hasattr(errno, errcode)
        ):
            errcode = "EINVAL"
        if errcode:
            errno_val = getattr(errno, errcode)
            raise OSError(errno_val, os.strerror(errno_val), orig_address)

class StubbedSSHAgentSocketWithAddressAndDeterministicDSA(
    StubbedSSHAgentSocketWithAddress
):
    """A [`StubbedSSHAgentSocketWithAddress`][] supporting deterministic DSA."""  # noqa: E501

def __init__(self) -> None:
        """Initialize the agent.

Set the supported extensions, and try issuing RFC 6979 and
        Pageant 0.68–0.80 DSA/ECDSA signatures, if possible.  See the
        [superclass constructor][StubbedSSHAgentSocketWithAddress] for
        other details.

Raises:
            KeyError: See superclass.
            OSError: See superclass.

"""  # noqa: RUF002
        super().__init__("query", "list-extended@putty.projects.tartarus.org")
        self.try_rfc6979 = True
        self.try_pageant_068_080 = True

STUB_AGENT_ENTRY_POINT = d_sasp.SSHAgentSocketProviderEntry(
    provider=StubbedSSHAgentSocket,
    key="stub_agent",
    aliases=(),
)
STUB_AGENT_WITH_ADDRESS_ENTRY_POINT = d_sasp.SSHAgentSocketProviderEntry(
    provider=StubbedSSHAgentSocketWithAddress,
    key="stub_agent_with_address",
    aliases=(),
)
STUB_AGENT_WITH_ADDRESS_AND_DETERMINISTIC_DSA_ENTRY_POINT = (
    d_sasp.SSHAgentSocketProviderEntry(
        provider=StubbedSSHAgentSocketWithAddressAndDeterministicDSA,
        key="stub_agent_with_address_and_deterministic_dsa",
        aliases=(),
    )
)

fakesshagent.git

Initial import of code from derivepassphrase