# SPDX-FileCopyrightText: 2026 Marco Ricci # # SPDX-License-Identifier: Zlib """Machinery for the fake SSH agent.""" from __future__ import annotations import collections import enum import errno import os import sched import struct import time from typing import TYPE_CHECKING, cast from typing_extensions import ( Literal, NamedTuple, TypeAlias, assert_never, overload, ) import derivepassphrase_sshagentsocketprovider as d_sasp from fakesshagent import _types, data, events if TYPE_CHECKING: import types from collections.abc import Callable, Generator, Iterable from collections.abc import Set as AbstractSet from typing_extensions import Buffer, Self Timefunc: TypeAlias = Callable[[], int | float] __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 InvalidStateError(ValueError): """The operation cannot continue because the system is not in the necessary state.""" # noqa: E501 class FrozenClockScheduler: """A scheduler that works with fixed timestamps.""" BLOCKING_CALL = "Blocking call in non-blocking scheduler?!" NON_MONOTONIC_TIMESTAMPS = "non-monotonic timestamps" def __init__(self, last_timestamp: float = float("-inf"), /) -> None: """Initialize the frozen clock scheduler. Args: last_timestamp: The initial timestamp to use. """ self.last_timestamp: float = last_timestamp """ The last timestamp value seen. Used to ensure monotonic progression of the timestamp. """ self.scheduler = sched.scheduler( timefunc=lambda: self.last_timestamp, delayfunc=self._do_not_call ) """The underlying scheduler.""" def run(self, timestamp: float, /) -> float | None: """Run the events scheduled at or before the timestamp. Args: timestamp: The timestamp to move forward to. Must not be smaller than the last timestamp, i.e., the clock must advance weakly monotonically. Returns: The distance from the given timestamp to the next scheduled event not handled in this call, if any. Raises: ValueError: The clock did not advance weakly monotonically. """ if timestamp < self.last_timestamp: # pragma: no cover [failsafe] raise ValueError(self.NON_MONOTONIC_TIMESTAMPS) self.last_timestamp = timestamp return self.scheduler.run(blocking=False) def _do_not_call( self, _ignored: float, / ) -> None: # pragma: no cover [failsafe] # Always raise an error. raise AssertionError(self.BLOCKING_CALL) class SSHAgentProtocolStackState(str, enum.Enum): pass class SSHAgentProtocolStack: """A Sans-I/O protocol stack for the SSH agent protocol. The protocol stack handles serialisation and deserialisation of protocol messages, and management of agent state. The stack translates these messages and state changes into events, such as [automatic server responses][events.ServerAction] and [agent actions][events.AgentAction]. The agent implementation (the "user" of this protocol stack) then interacts with the protocol stack in terms of these events, in a request-response manner. I/O from a client connection is represented by a [`BytesFromClient` agent input event][events.BytesFromClient], I/O to a client connection by a [`BytesToClient` agent action event][events.BytesToClient]; other [protocol stack feedback][events.ProtocolStackFeedback]/[agent actions][events.AgentAction], [reactions][events.AgentReactions] and [input][events.AgentInput] are also possible. The protocol stack itself contains no cryptographic code. In particular, **the protocol stack cannot fulfill signature requests and cannot securely store key material**. The agent implementation is responsible for both of these functionalities, by physically reacting to the appropriate agent actions and passing the correct agent reactions back. The protocol stack also manages key constraints automatically, including key timeouts, by signaling appropriate agent actions to the agent implementation. However, the protocol stack does not block, sleep or otherwise interface with the system clock and thread or process suspension functionality. Therefore, **the protocol stack cannot inform the agent implementation of key expirations in a timely manner unless the agent implementation regularly updates the protocol stack's knowledge of the current time**, at whatever frequency the implementation deems sufficient. The protocol stack uses handles to represent the I/O channels to different clients. See the [`Handle`][] class for more information. """ class Handle: """A handle for an SSH agent protocol stack. A handle represents an endpoint for feeding agent events into the protocol stack, and receiving agent actions from it. I/O to and from a client connection is mediated by the agent implementation, and expressed by the [`BytesFromClient` agent input event][events.BytesFromClient] and the [`BytesToClient` agent action][events.BytesToClient]. Agent implementations SHOULD create separate handles for each distinct client network connection they do I/O with. All handles from the same protocol stack share the same state, except for those endpoints. This includes the keystore in the agent implementation: key management actions are global to the protocol stack, may be triggered without client interaction, and thus may be returned as part of the next "list of events" response to the agent implementation (whatever handle that turns out to belong to). Handles are resources, and SHOULD be used as context managers to avoid lost messages. When entered and exited, the handle ensures that all input and output channels are empty. Once exited, a handle cannot be meaningfully interacted with anymore. """ _ALREADY_CLOSED = "the handle is already closed" """A common error message.""" _QUEUES_NOT_EMPTY = "the queues are not empty" """A common error message.""" def __init__(self, protocol_stack: SSHAgentProtocolStack, /) -> None: """Initialize a handle. Args: protocol_stack: The SSH agent protocol stack this handle belongs to. """ self.protocol_stack = protocol_stack """The protocol stack this handle belongs to.""" self.bytes_queue = bytearray() """ The queue of bytes from the client, containing not-yet-processed as well as partial requests. It is [fed][feed] from [`BytesFromClient` agent input events][events.BytesFromClient]. """ self.closed = False """True if and only if the handle is closed.""" def __enter__(self, /) -> Self: """Enter the handle as a new context. As part of this, check that the handle is not yet closed, and that the bytes queue from the client is empty. Returns: The handle. Raises: InvalidStateError: The handle is closed, or the bytes queue is not empty. """ if self.closed: raise InvalidStateError(self._ALREADY_CLOSED) if self.bytes_queue: raise InvalidStateError(self._QUEUES_NOT_EMPTY) return self def __exit__( self, _exc_type: type[BaseException] | None, _exc_value: BaseException | None, _exc_tb: types.TracebackType | None, /, ) -> bool: """Close the handle, and check that the byte queue is empty. Returns: False. This context manager does not swallow any exceptions. Raises: InvalidStateError: The bytes queue is not empty. """ self.closed = True if self.bytes_queue: raise InvalidStateError(self._QUEUES_NOT_EMPTY) return False @overload def feed( self, event: events.AgentInput, /, *, automatic_events: bool = False, ) -> tuple[events.ProtocolStackFeedback, ...]: ... @overload def feed( self, event: events.AgentEvent, /, *, automatic_events: Literal[False] = False, ) -> tuple[events.ProtocolStackFeedback, ...]: ... @overload def feed( self, event: events.AgentEvent, /, *, automatic_events: Literal[True] = True, ) -> tuple[ events.ProtocolStackFeedback | events.ServerEvent, ..., ]: ... def feed( self, event: events.AgentEvent, /, *, automatic_events: bool = False, ) -> tuple[ events.ProtocolStackFeedback | events.ServerEvent, ..., ]: """Feed agent events to the handle. Receive agent actions in return. If requested, also receive server events. Args: event: The agent event to process in the protocol stack. Client input is represented by a [`BytesFromClient` agent input event][events.BytesFromClient]. automatic_events: If true, then also receive automatic server events: the [parsed client requests][events.ClientRequest], the [server responses][events.ServerResponse], and the [automatic server actions][events.ServerAction]. Returns: A list of agent actions to perform, as a result of the processed agent event. May include a [`BytesToClient` agent action][events.BytesToClient], indicating which bytes to send to the client. If `automatic_events` is true, and the processed agent event is an [agent reaction][events.AgentReaction], then the list may also contain parsed client requests, server responses and automatic server actions. """ raise NotImplementedError def _process_event( self, request: events.ClientRequest, / ) -> tuple[events.ProtocolStackFeedback | events.ServerEvent, ...]: raise NotImplementedError class _KeyInfo(NamedTuple): """Information for key material managed by the agent. Attributes: constraints: The set of [key constraints][events.KeyConstraint] this key is governed by. token_id: The (opaque) identifier of the smartcard token the key is located on, if any. Otherwise, `None`. expiry_event: The current [scheduler event][sched.scheduler.queue] for the key's expiry, if the key expires. Otherwise, `None`. """ constraints: AbstractSet[events.KeyConstraint] """""" token_id: str | None """""" expiry_event: sched.Event | None """""" def __init__( self, /, *, timefunc: Timefunc | None = time.monotonic ) -> None: """Initialise the protocol stack. Args: timefunc: An optional zero-argument callable that returns the "current time", so that handles can automatically generate [`TimeProgressed` events][events.TimeProgressed] to inform the protocol stack of the passage of time. Defaults to [`time.monotonic`][]. If explicitly set to `None`, then no `TimeProgressed` events are automatically generated, and the agent implementation needs to update the protocol stack's clock itself, lest it remain static. """ self.timefunc = timefunc """ An optional zero-argument callable that returns the "current time". If given, then handles automatically generate [`TimeProgressed` events][events.TimeProgressed] to inform the protocol stack of the passage of time at the beginning of each [event cycle][feed]. The protocol stack's clock remains static between updates. """ self.scheduler = FrozenClockScheduler( timefunc() if timefunc is not None else float("-inf") ) """ An event scheduler for keeping track of key expiry events. Only ever run in non-blocking mode, and with a static clock set to the current value of [timestamp][]. """ self.recently_expired: collections.deque[bytes] = collections.deque() """ A queue of recently expired keys that have not yet been signaled to the agent implementation. Filled by the [scheduler][], and drained by [polling][poll]. """ self.known_keys: dict[bytes, SSHAgentProtocolStack._KeyInfo] = {} """ A dict of public keys to [key info][SSHAgentProtocolStack._KeyInfo]. """ def handle(self) -> Handle: """Return a newly constructed handle.""" return self.Handle(self) def poll( self, agent_input: events.AgentInput, / ) -> tuple[events.ProtocolStackFeedback, ...]: if isinstance(agent_input, events.TimeProgressed): next_timeout_delta = self.scheduler.run(agent_input.timestamp) expired_keys = self.recently_expired self.recently_expired = collections.deque() result: collections.deque[events.ProtocolStackFeedback] = ( collections.deque() ) for key in expired_keys: try: key_info = self.known_keys.pop(key) except KeyError as e: msg = f"Invalid known_keys state, missing {key!r}" raise AssertionError(msg) from e result.append( events.PurgeKey( public_key=key, token_id=key_info.token_id, client_initiated=False, ) ) if next_timeout_delta is not None: next_timeout = agent_input.timestamp + next_timeout_delta result.appendleft(events.NextTimedEventAt(next_timeout)) return tuple(result) assert_never(agent_input) 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:` 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=(), ) )