Split SSHAgentClient.request in two
Marco Ricci commited on 2026-06-13 16:04:08
Zeige 1 geänderte Dateien mit 50 Einfügungen und 19 Löschungen.
Factor out the serialization of the request message and the deserialization of the response message from `SSHAgentClient.request` into separate helper methods. `SSHAgentClient.request` now only contains the management of expected response codes, and the raising of errors if unexpected response codes are received. While this refactoring was unnecessary so far – requests and responses were always paired and aligned – the next piece of code I intend to commit requires handling unpaired responses, and will use only one of these helper methods.
- src/derivepassphrase/ssh_agent/__init__.py 3d1a7cc..0b83485
| ... | ... |
@@ -462,6 +462,50 @@ class SSHAgentClient: |
| 462 | 462 |
v() for v in known_good_agents.values() |
| 463 | 463 |
) |
| 464 | 464 |
|
| 465 |
+ def _send_request_message( |
|
| 466 |
+ self, code: int | _types.SSH_AGENTC, payload: Buffer, / |
|
| 467 |
+ ) -> None: |
|
| 468 |
+ """Prepare and send a generic request to the SSH agent. |
|
| 469 |
+ |
|
| 470 |
+ We take care of packing, framing and sending the request. |
|
| 471 |
+ |
|
| 472 |
+ Args: |
|
| 473 |
+ code: |
|
| 474 |
+ See [request][]. |
|
| 475 |
+ payload: |
|
| 476 |
+ See [request][]. |
|
| 477 |
+ |
|
| 478 |
+ """ |
|
| 479 |
+ payload = memoryview(payload) |
|
| 480 |
+ request_message = bytearray([code]) |
|
| 481 |
+ request_message.extend(payload) |
|
| 482 |
+ self._connection.sendall(self.string(request_message)) |
|
| 483 |
+ |
|
| 484 |
+ def _get_response_pair(self) -> tuple[int, bytes]: |
|
| 485 |
+ """Read the response for a generic request to the SSH agent. |
|
| 486 |
+ |
|
| 487 |
+ Returns: |
|
| 488 |
+ A 2-tuple consisting of the response code and the payload, |
|
| 489 |
+ with all wire framing removed. |
|
| 490 |
+ |
|
| 491 |
+ Raises: |
|
| 492 |
+ EOFError: |
|
| 493 |
+ The response from the SSH agent is truncated or missing. |
|
| 494 |
+ OSError: |
|
| 495 |
+ There was a communication error with the SSH agent. |
|
| 496 |
+ |
|
| 497 |
+ """ |
|
| 498 |
+ chunk = self._connection.recv(HEAD_LEN) |
|
| 499 |
+ if len(chunk) < HEAD_LEN: |
|
| 500 |
+ msg = "cannot read response length" |
|
| 501 |
+ raise EOFError(msg) |
|
| 502 |
+ response_length = int.from_bytes(chunk, "big", signed=False) |
|
| 503 |
+ response = self._connection.recv(response_length) |
|
| 504 |
+ if len(response) < response_length: |
|
| 505 |
+ msg = "truncated response from SSH agent" |
|
| 506 |
+ raise EOFError(msg) |
|
| 507 |
+ return response[0], response[1:] |
|
| 508 |
+ |
|
| 465 | 509 |
@overload |
| 466 | 510 |
def request( |
| 467 | 511 |
self, |
| ... | ... |
@@ -541,26 +585,13 @@ class SSHAgentClient: |
| 541 | 585 |
""" |
| 542 | 586 |
if isinstance(response_code, int): # pragma: no branch |
| 543 | 587 |
response_code = frozenset({response_code})
|
| 544 |
- payload = memoryview(payload) |
|
| 545 |
- request_message = bytearray([ |
|
| 546 |
- code if isinstance(code, int) else code.value |
|
| 547 |
- ]) |
|
| 548 |
- request_message.extend(payload) |
|
| 549 |
- self._connection.sendall(self.string(request_message)) |
|
| 550 |
- chunk = self._connection.recv(HEAD_LEN) |
|
| 551 |
- if len(chunk) < HEAD_LEN: |
|
| 552 |
- msg = "cannot read response length" |
|
| 553 |
- raise EOFError(msg) |
|
| 554 |
- response_length = int.from_bytes(chunk, "big", signed=False) |
|
| 555 |
- response = self._connection.recv(response_length) |
|
| 556 |
- if len(response) < response_length: |
|
| 557 |
- msg = "truncated response from SSH agent" |
|
| 558 |
- raise EOFError(msg) |
|
| 588 |
+ self._send_request_message(code, payload) |
|
| 589 |
+ actual_code, response = self._get_response_pair() |
|
| 559 | 590 |
if not response_code: # pragma: no cover [failsafe] |
| 560 |
- return response[0], response[1:] |
|
| 561 |
- if response[0] not in response_code: |
|
| 562 |
- raise SSHAgentFailedError(response[0], response[1:]) |
|
| 563 |
- return response[1:] |
|
| 591 |
+ return actual_code, response |
|
| 592 |
+ if actual_code not in response_code: |
|
| 593 |
+ raise SSHAgentFailedError(actual_code, response) |
|
| 594 |
+ return response |
|
| 564 | 595 |
|
| 565 | 596 |
def list_keys(self) -> Sequence[_types.SSHKeyCommentPair]: |
| 566 | 597 |
"""Request a list of keys known to the SSH agent. |
| 567 | 598 |